axon
axon copied to clipboard
Optionally return gradients / gradient state in train step