Added error handling for the VAE state_dict, in some occasions we mig…

…ht end up with it throwing a RuntimeError, in this case we fallback to using strict=False when loading the state_dict.

muse_maskgit_pytorch/vqgan_vae.py

@@ -439,7 +439,11 @@ def state_dict(self, *args, **kwargs):
 
     @remove_vgg
     def load_state_dict(self, *args, **kwargs):
-        return super().load_state_dict(*args, **kwargs)
+        try:
+            return super().load_state_dict(*args, **kwargs)
+        except RuntimeError:
+            return super().load_state_dict(*args, **kwargs, strict=False)
+
 
     def save(self, path):
         if self.accelerator is not None: