num_epochs = 10
for epoch in range(num_epochs):
cum_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.as_in_context(ctx)
# ctx can be mx.gpu() if gpu available otherwise mx.cpu()
labels = labels.as_in_context(ctx)
# ctx can be mx.gpu() if gpu available otherwise mx.cpu()
with autograd.record():
outputs = net(inputs)
loss = loss_function(outputs, labels)
# here you are computing the loss
loss.backward()
# here you are differentiating your loss with respect to all your model parameters.
trainer.step(batch_size=inputs.shape[0])
# running trainer
cum_loss += loss.mean()
# accumulating loss of current batch
epoch_loss = cum_loss.asscalar()/len(train_loader)
# len(train_loader) gives number of batches, so you are dividing all accumulated loss by num of batches
print('After epoch {}: Loss: {}'.format(epoch + 1, epoch_loss))