criterion=smp.utils.losses.DiceLoss()
for epoch in range(epochs):
        model.train()
        epoch_loss=0
        for i,(image,mask) in enumerate(train_dl):
                 data=torch.randn(8,3,128,128)
                 target=torch.randn(8,4,128,128)
                 optimizer.zero_grad()
                 data=data.to(device).float()
                 output=model(data)
                 output=torch.sigmoid(output)
                 output=output.to('cpu')
                 loss=criterion(output,target)
                 epoch_loss += loss
                 loss.backward()
                 optimizer.step()
                 scheduler.step()
           
        epoch_loss = epoch_loss / total_steps
        print('epoch:{},epoch_loss:{}'.format(epoch,epoch_loss))

I’m sorry for my slow response。This code might suit your needs。
This is the error：

捕获1114×152 7.7 KB

You can see that the first call to scheduler. Step should have been fault-free because it printed the following print statement, as well as eval for the model. But an error should have occurred on the second call