Update prediction.py

prediction.py

@@ -1,5 +1,6 @@
 import os
 os.chdir('..')
+base_dir = os.getcwd()
 from dataloader import CellLoader
 from celle_main import instantiate_from_config
 from omegaconf import OmegaConf
@@ -46,7 +47,9 @@ def run_image_prediction(
     config["model"]["params"]["vqgan_model_path"] = None
 
     # Instantiate model from config and move to device
-    model = instantiate_from_config(config).to(device)
+    model = instantiate_from_config(config.model).to(device)
+
+    os.chdir(base_dir)
 
     # Sample from model using provided sequence and nucleus image
     _, _, _, predicted_threshold, predicted_heatmap = model.celle.sample(