save or load checkpoint

[rancheng]

distributed_tutorial/src/mnist-mixed.py

Lines 108 to 114 in 2446796

	print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
	epoch + 1,
	args.epochs,
	i + 1,
	total_step,
	loss.item())
	)

could you add the save model line to the example to be more complete. Thanks!

torch.save(model.state_dict(), CHECKPOINT_PATH)