Spaces:

rayli
/

DragAPart

Running on Zero

Ruining Li commited on Mar 16

Commit

effdf42

•

1 Parent(s): 50e01f3

Minor change.

Files changed (1) hide show

app.py CHANGED Viewed

@@ -89,7 +89,7 @@ def model_init():
         ),
         load=False,
     )
-    model.load_state_dict(torch.load(model_checkpoint)["model"])
     model = model.to("cuda")
     return model

         ),
         load=False,
     )
+    model.load_state_dict(torch.load(model_checkpoint, map_location="cpu")["model"])
     model = model.to("cuda")
     return model