From dcd8ab567231f43e6bf627ba4b6584994dfd6d19 Mon Sep 17 00:00:00 2001 From: Jonathan Ventura Date: Wed, 22 Jun 2022 09:57:42 -0700 Subject: [PATCH] Only save network_fine state dict if it exists --- run_nerf.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/run_nerf.py b/run_nerf.py index bc270be86..57fbc79a8 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -791,12 +791,14 @@ def train(): # Rest is logging if i%args.i_weights==0: path = os.path.join(basedir, expname, '{:06d}.tar'.format(i)) - torch.save({ + ckpt = { 'global_step': global_step, 'network_fn_state_dict': render_kwargs_train['network_fn'].state_dict(), - 'network_fine_state_dict': render_kwargs_train['network_fine'].state_dict(), 'optimizer_state_dict': optimizer.state_dict(), - }, path) + } + if render_kwargs_train['network_fine'] is not None: + ckpt.update({'network_fine_state_dict': render_kwargs_train['network_fine'].state_dict()}) + torch.save(ckpt, path) print('Saved checkpoints at', path) if i%args.i_video==0 and i > 0: