Spaces:

SaniaE
/

GeoGen-API

Sleeping

saniaE commited on Apr 8

Commit

a510479

1 Parent(s): 2aac227

updated gen model load state

Files changed (1) hide show

app.py CHANGED Viewed

@@ -47,7 +47,7 @@ def load_model():
         checkpoint = torch.load(model_path, map_location=DEVICE)
         gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
-        gen_model.load_state_dict(checkpoint["gen_state_dict"])
         gen_model.eval()
         print("✅ SUCCESS: Petrol Pump GAN is live!")
     except Exception as e:

         checkpoint = torch.load(model_path, map_location=DEVICE)
         gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
+        gen_model.load_state_dict(checkpoint["gen_state_dict"], strict=False)
         gen_model.eval()
         print("✅ SUCCESS: Petrol Pump GAN is live!")
     except Exception as e: