Update consistencytta.py

consistencytta.py

@@ -75,9 +75,17 @@ class ConsistencyTTA(nn.Module):
             [self.text_encoder, self.vae, self.fn_STFT, self.unet],
             ['text_encoder', 'vae', 'fn_STFT', 'unet']
         ):
-            assert model.training == False, f"The {name} is not in eval mode."
+            try:
+                assert model.training == False, f"The {name} is not in eval mode."
+            except:
+                model.eval()
+                assert model.training == False, f"The {name} is not in eval mode."
             for param in model.parameters():
-                assert param.requires_grad == False, f"The {name} is not frozen."
+                try:
+                    assert param.requires_grad == False, f"The {name} is not frozen."
+                except:
+                    param.requires_grad_(False)
+                    assert param.requires_grad == False, f"The {name} is not frozen."
 
 
     @torch.no_grad()