The error still persists with your change; this is what my training loop looks like now:
model.to(device)

arr = []
models = []
epochs = 10
def train():
    criterion = DiceLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    num_steps_train = len(train_loader)
    
    print(num_steps_train)
    
    for epoch in range(epochs):
        t0 = time.time()
        for i, (images, masks) in enumerate(train_loader):
            images = images.to(device)
            masks = masks.to(device)
            outputs = model(images)
            outputs = torch.argmax(outputs, 1, keepdim=True).float()
            loss = criterion(outputs, masks)
            arr.append(loss)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print('\r Epoch {} of {}   {:.2f}% - loss: {:.4f} - ({:.2f}s)'
                       .format(epoch + 1, epochs, i * 100 / num_steps_train, float
                               (loss), time.time() - t0))
        models.append(model)

here is the full stack trace when i run train()
10
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-14-2da0ffaf5447> in <module>()
----> 1 train()

2 frames
<ipython-input-13-9df18eabb960> in train()
     21             arr.append(loss)
     22             optimizer.zero_grad()
---> 23             loss.backward()
     24             optimizer.step()
     25             print('\r Epoch {} of {}   {:.2f}% - loss: {:.4f} - ({:.2f}s)'

/usr/local/lib/python3.6/dist-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    183                 products. Defaults to ``False``.
    184         """
--> 185         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    186 
    187     def register_hook(self, hook):

/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    125     Variable._execution_engine.run_backward(
    126         tensors, grad_tensors, retain_graph, create_graph,
--> 127         allow_unreachable=True)  # allow_unreachable flag
    128 
    129 

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn