Update loader.py

loader.py

@@ -31,9 +31,9 @@ def load_graph_decoder(path='model_labeled'):
     model.init_model(path)
     model.disable_grads()
 
-    # trainable_params, all_param = count_parameters(model)
-    # param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
-    #     path, trainable_params, all_param, 100 * trainable_params / all_param
-    # )
-    # print(param_stats)
+    trainable_params, all_param = count_parameters(model)
+    param_stats = "Loaded Graph DiT from {} trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
+        path, trainable_params, all_param, 100 * trainable_params / all_param
+    )
+    print(param_stats)
     return model