@ptrblck  Below is my training loop.
with torch.enable_grad():
classifier = classifier.train()
        for i, (points, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
            optimizer.zero_grad()

            if i == 6862:
                break

            points = points.data.numpy()
            target = preprocessed(target)
            # points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3])
            points = torch.Tensor(points)
            points, target = points.float().cuda(), target.cuda()
            points = torch.reshape(points, (BATCH_SIZE, 36, NUM_POINT, 3))
            points = points.transpose(3, 2)
            # target1 =target

            seg_pred = classifier(points)
            print(seg_pred.shape)

            seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
            seg_pred =torch.argmax(seg_pred,1)
            # print(seg_pred.shape)


            target = target.view(-1, 1)[:, 0]
            print(target.shape)

            loss = criterion(seg_pred, target.cuda())

            seg_pred =seg_pred.cpu().detach().numpy()
            # target1 = target.clone().detach()
            # seg_pred1 = seg_pred.clone().detach()


            # loss = lovasz_softmax(seg_pred, target, ignore=None) + loss_func(seg_pred, target)
            print(loss)
            # loss.backward()
            #

            loss.backward()
            optimizer.step()

IoU loss I am using as below:
class IoULoss(torch.nn.Module):
def init(self, weight=None, size_average=True):
super(IoULoss, self).init()
def forward(self, inputs, targets, smooth=1):
    import torch.nn.functional as F
    # comment out if your model contains a sigmoid or equivalent activation layer
    inputs = F.sigmoid(inputs)

    # flatten label and prediction tensors
    inputs = inputs.view(-1)
    targets = targets.view(-1)

    # intersection is equivalent to True Positive count
    # union is the mutually inclusive area of all labels & predictions
    intersection = (inputs * targets).sum()
    total = (inputs + targets).sum()
    union = total - intersection

    IoU = (intersection + smooth) / (union + smooth)

    return 1 - IoU

classifier = MODEL.segmodel(NUM_CLASSES).cuda()
# criterion = MODEL.get_loss().cuda()
criterion = IoULoss().cuda()
classifier.apply(inplace_relu)