I want to do a binary boundary losses inspired from https://github.com/JunMa11/SegLoss/blob/master/losses_pytorch/boundary_loss.py 2 and https://arxiv.org/abs/1812.07032. I have seen the use of with torch.no_grad in  forward method in loss. I am not sure how torch.no_grad is recommended in the forward method. Since it is done only in the labels which are not yet attach to the computational graph does it is a good idea to use it. Will it influence training.
Is better to have multiple inputs like def forward(self, preds, targets, distance_map):?
"""
Inspired from https://github.com/JunMa11/SegLoss/blob/master/losses_pytorch/boundary_loss.py.
Which list a lot of pytorch losses for segmentatoin.
"""

import numpy as np
from scipy.ndimage import distance_transform_edt
import torch


def compute_edts_forhdloss(segmentation):
    res = np.zeros(segmentation.shape)
    for i in range(segmentation.shape[0]):
        posmask = segmentation[i]
        negmask = ~posmask
        res[i] = distance_transform_edt(posmask) + distance_transform_edt(negmask)
    return res


class BinaryBDLoss(torch.nn.Module):
    def __init__(self, num_channel: int, ndim=4, per_channel: bool = True, reduction: str = "mean"):
        """
        compute boudary loss
        only compute the loss of foreground
        ref: https://github.com/LIVIAETS/surface-loss/blob/108bd9892adca476e6cdf424124bc6268707498e/losses.py#L74
        """
        super().__init__()
        assert ndim in [4, 5]
        self.per_channel = per_channel
        self.num_channel = num_channel
        self.ndim = ndim
        self.axis = list(range(ndim))
        self.reduction = reduction

        if self.num_channel > 1 & self.per_channel:
            if self.ndim == 5:
                self.axis = [0, -3, -2, -1]
            if self.ndim == 4:
                self.axis = [0, -2, -1]
        # self.do_bg = do_bg

    def forward(self, preds, targets):
        """
        preds: (batch_size, class, x,y)
        target: ground truth, shape: (batch_size, 1, x,y)
        bound: precomputed distance map, shape (batch_size, class, x,y)
        """
        if preds.shape[1] != self.num_channel:
            preds = preds.view(preds.shape[0], -1, *preds.shape[1:])
            targets = targets.view(*preds.shape)
            # distance = distance.view(*preds.shape)

        with torch.no_grad():
            bound = compute_edts_forhdloss(targets.cpu().numpy() > 0.5)

        bound = torch.from_numpy(bound)
        bound.requires_grad = True
        bound.to(preds.device)
        bound.type(preds.dtype)

        # print('preds shape: ', preds.shape)
        # Self.idc is used to filter out some classes of the target mask. Use fancy indexing
        # self.idc: List[int] = kwargs["idc"]
        # pc = preds[:, self.idc, ...].type(torch.float32)
        # dc = bound[:, self.idc, ...].type(torch.float32)

        if preds.ndim == 5:
            multipled = torch.einsum("bcxyz,bcxyz->bcxyz", preds, bound)
        else:
            multipled = torch.einsum("bcwh,bcwh->bcwh", preds, bound)

        bd_loss = torch.mean(multipled, dim=self.axis)

        if self.reduction == 'mean':
            bd_loss = bd_loss.mean()
            return bd_loss

        if self.reduction == 'sum':
            bd_loss = bd_loss.sum()
            return bd_loss

        if self.reduction == 'none':
            return bd_loss

        raise Exception('Unexpected reduction {}'.format(self.reduction))