Added optimizer.zero_grad(), loss.backward(), and optimizer.step() properly 6eb22ad verified Bhavibond commited on Mar 20