Hi, I am having issues with Dice Loss and Pytorch Ignite. I am trying to reproduce the result of Ternausnet using dice loss but my gradients keep being zero and loss just does not improve or shows very strange results (negative, nan, etc). I am not sure where to look for a possible source of the issue. Below is the code for DiceLoss:
from torch import nn
from torch.nn import functional as F

# from catalyst.contrib.nn import DiceLoss
import torch


class DiceLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, logits, targets, eps=1, threshold=None):

        # comment out if your model contains a sigmoid or
        # equivalent activation layer
        proba = torch.sigmoid(logits)
        proba = proba.view(proba.shape[0], 1, -1)
        targets = targets.view(targets.shape[0], 1, -1)
        if threshold:
            proba = (proba > threshold).float()
        # flatten label and prediction tensors

        intersection = torch.sum(proba * targets, dim=1)
        summation = torch.sum(proba, dim=1) + torch.sum(targets, dim=1)
        dice = (2.0 * intersection + eps) / (summation + eps)
        # print(intersection, summation, dice)
        return (1 - dice).mean()


and here is the model (unet11 backbone code taken from ternausnet):
import pytorch_lightning as pl
import torch
from torch import nn
from torchvision import models
from carvana_unet.utils import DiceLoss

def conv3x3(in_: int, out: int) -> nn.Module:
    return nn.Conv2d(in_, out, 3, padding=1)


class ConvRelu(nn.Module):
    def __init__(self, in_: int, out: int) -> None:
        super().__init__()
        self.conv = conv3x3(in_, out)
        self.activation = nn.ReLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        x = self.activation(x)
        return x


class DecoderBlock(nn.Module):
    def __init__(
        self, in_channels: int, middle_channels: int, out_channels: int
    ) -> None:
        super().__init__()

        self.block = nn.Sequential(
            ConvRelu(in_channels, middle_channels),
            nn.ConvTranspose2d(
                middle_channels,
                out_channels,
                kernel_size=3,
                stride=2,
                padding=1,
                output_padding=1,
            ),
            nn.ReLU(inplace=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.block(x)


class UNet11Lightning(pl.LightningModule):
    def __init__(
        self, num_filters: int = 32, pretrained: bool = True, loss_fn=DiceLoss
    ) -> None:
        """
        Args:
            num_filters:
            pretrained:
                False - no pre-trained network is used
                True  - encoder is pre-trained with VGG11
        """
        super().__init__()
        self.loss_fn = loss_fn()
        self.pool = nn.MaxPool2d(2, 2)

        self.encoder = models.vgg11(pretrained=pretrained).features

        self.relu = self.encoder[1]
        self.conv1 = self.encoder[0]
        self.conv2 = self.encoder[3]
        self.conv3s = self.encoder[6]
        self.conv3 = self.encoder[8]
        self.conv4s = self.encoder[11]
        self.conv4 = self.encoder[13]
        self.conv5s = self.encoder[16]
        self.conv5 = self.encoder[18]

        self.center = DecoderBlock(
            num_filters * 8 * 2, num_filters * 8 * 2, num_filters * 8
        )
        self.dec5 = DecoderBlock(
            num_filters * (16 + 8), num_filters * 8 * 2, num_filters * 8
        )
        self.dec4 = DecoderBlock(
            num_filters * (16 + 8), num_filters * 8 * 2, num_filters * 4
        )
        self.dec3 = DecoderBlock(
            num_filters * (8 + 4), num_filters * 4 * 2, num_filters * 2
        )
        self.dec2 = DecoderBlock(
            num_filters * (4 + 2), num_filters * 2 * 2, num_filters
        )
        self.dec1 = ConvRelu(num_filters * (2 + 1), num_filters)

        self.final = nn.Conv2d(num_filters, 1, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv1 = self.relu(self.conv1(x))
        conv2 = self.relu(self.conv2(self.pool(conv1)))
        conv3s = self.relu(self.conv3s(self.pool(conv2)))
        conv3 = self.relu(self.conv3(conv3s))
        conv4s = self.relu(self.conv4s(self.pool(conv3)))
        conv4 = self.relu(self.conv4(conv4s))
        conv5s = self.relu(self.conv5s(self.pool(conv4)))
        conv5 = self.relu(self.conv5(conv5s))

        center = self.center(self.pool(conv5))

        dec5 = self.dec5(torch.cat([center, conv5], 1))
        dec4 = self.dec4(torch.cat([dec5, conv4], 1))
        dec3 = self.dec3(torch.cat([dec4, conv3], 1))
        dec2 = self.dec2(torch.cat([dec3, conv2], 1))
        dec1 = self.dec1(torch.cat([dec2, conv1], 1))
        return self.final(dec1)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def training_step(self, batch, batch_idx):
        x, y = batch["features"].cuda(), batch["target"].cuda()
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        result = pl.TrainResult(loss)
        result.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        x, y = batch["features"].cuda(), batch["target"].cuda()
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        result = pl.EvalResult(checkpoint_on=loss)
        result.log("val_loss", loss, prog_bar=True, on_step=True)
        tensorboard_log = {"val_loss": loss}
        return {"loss": loss, "log": tensorboard_log}

    def test_step(self, batch, batch_idx):
        x = batch["features"].cuda()
        y_hat = self(x)
        return torch.nn.functional.sigmoid(y_hat)


can anybody help me find the source of the issue or point me into the direction of where to look? I have tried other losses (like BCE with logits which was decreasing into -inf) and nothing seems to work.