Spaces:

karanravindra
/

DigitDreamer

Runtime error

karanravindra commited on 10 days ago

Commit

9f1fbd5

•

1 Parent(s): f6a41bd

change weights map location

Files changed (1) hide show

app.py CHANGED Viewed

@@ -18,8 +18,8 @@ torch.set_grad_enabled(False)
 decoder = Autoencoder().decoder
 dit = DiT()
-decoder.load_state_dict(torch.load("models/ft-decoder.pth", weights_only=True))
-dit.load_state_dict(torch.load("models/diffusion.pth", weights_only=True))
 rf = RF(dit)

 decoder = Autoencoder().decoder
 dit = DiT()
+decoder.load_state_dict(torch.load("models/ft-decoder.pth", weights_only=True, map_location="cpu"))
+dit.load_state_dict(torch.load("models/diffusion.pth", weights_only=True, map_location="cpu"))
 rf = RF(dit)