From 83a30e8a9d59aafd94faed000c855fb73fd361f1 Mon Sep 17 00:00:00 2001
From: liuwenran <448073814@qq.com>
Date: Fri, 8 Sep 2023 13:28:22 +0800
Subject: [PATCH] fix test

---
 mmagic/models/archs/wrapper.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/mmagic/models/archs/wrapper.py b/mmagic/models/archs/wrapper.py
index 5b95a9d649..b985ac231d 100644
--- 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