irreps_in = o3.Irreps('1x0e')   # Single scalars
lmax = 6
irreps_out = o3.Irreps.spherical_harmonics(lmax)  # Predict vectors

r_max = 1.5

model_kwargs = {
    'irreps_in': irreps_in, 'irreps_out': irreps_out,         # Data-types of input and output
    'max_radius': r_max, 'num_neighbors': 3, 'num_nodes': 6,  # Cutoff radius and numbers used to normalize conv
    'mul': 1, 'layers': 2, 'lmax': lmax, 'pool_nodes': True,  # Network details
}

# Create three random models
model1 = SimpleNetwork(**model_kwargs)
model2 = SimpleNetwork(**model_kwargs)
model3 = SimpleNetwork(**model_kwargs)