def forward_fn(inputs, targets):
    logits = model(inputs)
    loss = criterion(logits, targets)
    return loss, logits

# get grad function
grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

# define train step function
def train_step(inputs, targets):
    (loss, logits), grads = grad_fn(inputs, targets) # get values and gradients
    print('grads=',grads)
    optimizer(grads) # update gradient
    return loss, logits