update

nsr/lsgm/train_util_diffusion_lsgm_noD_joint.py

@@ -154,8 +154,11 @@ class SDETrainLoopJoint(TrainLoopDiffusionWithRec):
         self.ddp_rec_model = functools.partial(self.model, model_name='rec')
         self.ddp_ddpm_model = functools.partial(self.model, model_name='ddpm')
 
-        self.rec_model = self.ddp_model.module.rec_model
-        self.ddpm_model = self.ddp_model.module.ddpm_model  # compatability
+        # self.rec_model = self.ddp_model.module.rec_model
+        # self.ddpm_model = self.ddp_model.module.ddpm_model  # compatability
+
+        self.rec_model = self.ddp_model.rec_model
+        self.ddpm_model = self.ddp_model.ddpm_model  # compatability
 
         # TODO, required?
         # for param in self.ddp_rec_model.module.decoder.triplane_decoder.parameters(  # type: ignore