state = train.initilaize()
for _ in range(5):    
    train_data = _ #pick random clients    
    state, metrics = train.next(state, train_data)    
    print(metrics.loss)