fix SGD bug

2023-04-25 22:39:49 +08:00
parent dd1ecdd2a7
commit d4bf5f9984
5 changed files with 174 additions and 11 deletions
@@ -64,8 +64,8 @@ class fsModel(BaseModel):

        # Id network
        netArc_checkpoint = opt.Arc_path
-        netArc_checkpoint = torch.load(netArc_checkpoint)
-        self.netArc = netArc_checkpoint['model'].module
+        netArc_checkpoint = torch.load(netArc_checkpoint, map_location=torch.device("cpu"))
+        self.netArc = netArc_checkpoint
        self.netArc = self.netArc.to(device)
        self.netArc.eval()