Update usr/diff/shallow_diffusion_tts.py

usr/diff/shallow_diffusion_tts.py

@@ -231,6 +231,7 @@ class GaussianDiffusion(nn.Module):
     def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
                 ref_mels=None, f0=None, uv=None, energy=None, infer=False, **kwargs):
         b, *_, device = *txt_tokens.shape, txt_tokens.device
+        print(txt_tokens.device)
         ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
                        skip_decoder=(not infer), infer=infer, **kwargs)
         cond = ret['decoder_inp'].transpose(1, 2)