fix test

liuwenran · Sep 8, 2023 · 83a30e8 · 83a30e8
1 parent 850b7e3
commit 83a30e8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mmagic/models/archs/wrapper.py b/mmagic/models/archs/wrapper.py
@@ -183,5 +183,8 @@ def to(
         torch_device: Optional[Union[str, torch.device]] = None,
         torch_dtype: Optional[torch.dtype] = None,
     ):
-        self.model.to(torch_device, torch_dtype)
+        if torch_dtype is None:
+            self.model.to(torch_device)
+        else:
+            self.model.to(torch_device, torch_dtype)
         return self