Spaces:

HuangLab
/

CELL-E_2-Image_Prediction

Sleeping

Emaad commited on May 17, 2023

Commit

1d52ab8

1 Parent(s): 6b1f49f

Update celle_main.py

Files changed (1) hide show

celle_main.py CHANGED Viewed

@@ -198,7 +198,7 @@ class CELLE_trainer(pl.LightningModule):
                 if k.startswith(ik):
                     # print("Deleting key {} from state_dict.".format(k))
                     del ckpt[k]
-        self.load_state_dict(ckpt, strict=True)
         print(f"Restored from {path}")
     def forward(self, text, condition, target, return_loss=True):

                 if k.startswith(ik):
                     # print("Deleting key {} from state_dict.".format(k))
                     del ckpt[k]
+        self.load_state_dict(ckpt, strict=False)
         print(f"Restored from {path}")
     def forward(self, text, condition, target, return_loss=True):