Upload loader.py with huggingface_hub

loader.py

@@ -18,7 +18,7 @@ def count_parameters(model):
 
     return trainable_params, all_param
 
-def load_graph_decoder(device, path='model_labeled'):
+def load_graph_decoder(path='model_labeled'):
     model_config_path = f"{path}/config.yaml"
     data_info_path = f"{path}/data.meta.json"
 
@@ -30,8 +30,6 @@ def load_graph_decoder(device, path='model_labeled'):
     )
     model.init_model(path)
     model.disable_grads()
-    model.to(device)
-    print('Moving model to', device)
 
     trainable_params, all_param = count_parameters(model)
     param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(