fix: Runtime errors

main.py

@@ -22,6 +22,7 @@ def main():
                 data_root = args.data_root,
                 n_trials = 3 if args.debug else None,
                 size = 2**10 if args.debug else None,
+                method = args.method,
                 model_name=args.model,
                 model_path = args.model_load_path,
                 model_out = args.model_save_path
@@ -32,6 +33,8 @@ def main():
 
         case _:
             raise NotImplementedError(f"Mode {args.mode} is not implemented yet")
+    
+    print("Done")
 
 
 if __name__ == "__main__":