diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index f4ca94b..48fbdc3 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.pyr_32 import config + from configs.iccvw.pyr_64_cape_m import config else: config_module = importlib.import_module(args['config']) config = config_module.config @@ -549,7 +549,7 @@ def plot_everything(batch: dict, type: str): figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) } batch_id = batch['id'] - batch_frame_nr = batch['frame_nr'].cpu().numpy() + batch_frame_nr = batch['frame_nr'] batch_garment = batch['garment'] batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))] batch_info = text_to_html(batch_info) @@ -855,7 +855,7 @@ if __name__ == '__main__': optimizer = optim.Adam(model.parameters(), lr=config['lr']) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) - ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=5) scaler = torch.cuda.amp.GradScaler(enabled=use_amp) g_balancing_coeff = config['starting_balancing_coeff']