Update gaussian_diffusion.py

Author: mainpyp

src/modeling/diffusion/gaussian_diffusion.py

@@ -243,7 +243,7 @@ def training_losses(self, model, x_start_never_used, t, model_kwargs=None, noise
         """
         assert "input_ids" in model_kwargs
         input_ids = model_kwargs.pop("input_ids").to(t.device)
-        x_start_mean = model.model.get_embeds(input_ids)
+        x_start_mean = model.model.module.get_embeds(input_ids)
 
         std = _extract_into_tensor(
             self.sqrt_one_minus_alphas_cumprod,
@@ -256,7 +256,7 @@ def training_losses(self, model, x_start_never_used, t, model_kwargs=None, noise
         if noise is None:
             noise = th.randn_like(x_start)
         x_t = self.q_sample(x_start, t, noise=noise)  # reparametrization trick.
-        get_logits = model.model.get_logits
+        get_logits = model.model.module.get_logits
 
         terms = {}