can you try with this step-function:
def step(batch, model, criterion, optimizer=None):
    x = batch["X"].to(DEVICE)
    y = batch["y"].to(DEVICE)

    ## Forward Pass ##
    predictions = model(inputs)

    ## Calculate Loss ##
    loss = criterion(predictions, y)

    if optimizer is not None:
        optimizer.zero_grad()
        # backward pass + optimize
        loss.backward()
        optimizer.step()
        
    return loss

I don’t know whether it is a difference if you zero the grads by the optimizer’s function or the model’s function, but I usually do it with the optimizer and this works for me.