dodaltuin / passive-lv-gnn-emul Goto Github PK

# Here I have changed : sow_latents: bool = True
def __call__(self, V: jnp.ndarray, E: jnp.ndarray, theta: jnp.ndarray, z_global: jnp.ndarray = None, sow_latents: bool = True):
……….
# only need local representation for real nodes
z_local = z_local[self.real_node_indices]

# Here is my saving z_local code
z_local_np = np.array(z_local)
np.save('z_local_1.npy', z_local_np)

# save value of final learned representation if required for fixed geometry emulator
if sow_latents:
    return z_local