loading weights

app.py

@@ -9,7 +9,14 @@ import matplotlib.pyplot as plt
 import numpy as np
 import optax
 
-from dcmnet.data import prepare_batches
+
+import pandas as pd
+from dcmnet.modules import MessagePassingModel
+
+test_weights = pd.read_pickle("wbs/best_0.0_params.pkl")
+
+
+
 
 # Disable future warnings.
 import warnings
@@ -23,7 +30,9 @@ rotation = e3x.so3.random_rotation(rotation_key)
 
 st.write(rotation)
 
-st.write(prepare_batches)
+st.write(MessagePassingModel)
+st.write(test_weights)
+
 
 x = st.slider('Select a value')
 st.write(x, 'squared is', x * x)