Fix train.py for training (autograd + in-place ops)

85d3c87

verified

Reality123b commited on 17 days ago