Modify neural methods's `__call__` to do just one step

[michalk8]

Currently, the __call__ method takes in various dataloaders. In order to have more flexible way of training (e.g. for sharding), the __call__ should do just one step of the algorithm.
In a separate PR later, we will then define a generic trainer.