Element 0 of tensors does not require grad and does not have a grad_fn :RunTimeerror

Hello guys, I am new to this PyTorch environment :slightly_smiling_face:

Here, I鈥檓 trying to construct a new loss balancing technique for my Forward problem with PINNs using NVIDIA-modulus- nvidia / Modulus / Modulus 路 GitLab.

I would like to alter the losses during training step=0. For e.g., I would like to initialize each losses[key]=1 for step=0

class Aggregator(nn.Module):
    Base class for loss aggregators

    def __init__(self, params, num_losses, weights):
        self.params: List[torch.Tensor] = list(params)
        self.num_losses: int = num_losses
        self.weights: Optional[Dict[str, float]] = weights
        self.device: torch.device
        self.device = list(set(p.device for p in self.params))[0]
        self.init_loss: torch.Tensor = torch.tensor(0.0, device=self.device)

        def weigh_losses_initialize(
            weights: Optional[Dict[str, float]]
        ) -> Callable[
            [Dict[str, torch.Tensor], Optional[Dict[str, float]]],
            Dict[str, torch.Tensor],
            if weights is None:

                def weigh_losses(
                    losses: Dict[str, torch.Tensor], weights: None
                ) -> Dict[str, torch.Tensor]:
                    return losses


                def weigh_losses(
                    losses: Dict[str, torch.Tensor], weights: Dict[str, float]
                ) -> Dict[str, torch.Tensor]:
                    for key in losses.keys():
                        if key not in weights.keys():
                            weights.update({key: 1.0})
                    losses = {key: weights[key] * losses[key] for key in losses.keys()}
                    return losses

            return weigh_losses

        self.weigh_losses = weigh_losses_initialize(self.weights)

class SoftAdapt(Aggregator):
    SoftAdapt for loss aggregation
    Reference: "Heydari, A.A., Thompson, C.A. and Mehmood, A., 2019.
    Softadapt: Techniques for adaptive loss weighting of neural networks with multi-part loss functions.
    arXiv preprint arXiv: 1912.12355."

    def __init__(self, params, num_losses, eps=1e-8, weights=None):
        super().__init__(params, num_losses, weights)
        self.eps: float = eps
            "prev_losses", torch.zeros(self.num_losses, device=self.device)
            "lmbda_ema", torch.ones(self.num_losses, device=self.device)

    def forward(self, losses: Dict[str, torch.Tensor], step: int) -> torch.Tensor:
        Weights and aggregates the losses using the original variant of the softadapt algorithm

        losses : Dict[str, torch.Tensor]
            A dictionary of losses.
        step : int
            Optimizer step.

        loss : torch.Tensor
            Aggregated loss.

        # weigh losses
        losses = self.weigh_losses(losses, self.weights)

        # Initialize loss
        loss: torch.Tensor = torch.zeros_like(self.init_loss)

        # Aggregate losses by summation at step 0
        if step == 0:
            for i, key in enumerate(losses.keys()):
            self.lmbda_ema[i] = 1/(losses[key]+ self.eps)
            losses[key] = losses[key]* self.lmbda_ema[i]
            loss +=  losses[key].clone().detach()  # simply, initial losses[key] = 1

        # Aggregate losses using SoftAdapt for step > 0
            lmbda: torch.Tensor = torch.ones_like(self.prev_losses)
            lmbda_sum: torch.Tensor = torch.zeros_like(self.prev_loss)
            losses_stacked: torch.Tensor = torch.stack(list(losses.values()))
            normalizer: torch.Tensor = (losses_stacked / self.prev_losses).max()
            for i, key in enumerate(losses.keys()):
                with torch.no_grad():
                    lmbda[i] = torch.exp(
                        losses[key] / (self.prev_losses[i] + self.eps) - normalizer
                    lmbda_sum += lmbda[i]
                loss += lmbda[i].clone() * losses[key]
                self.prev_losses[i] = losses[key].clone().detach()
            loss *= self.num_losses / (lmbda_sum + self.eps)
        return loss

The error:

File "/opt/conda/lib/python3.8/site-packages/functorch/_src/monkey_patching.py", line 77, in _backward return _old_backward(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 402, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py", line 191, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass

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

Thanks in advance

You are explicitly detaching tensors in your custom function and are also running some forward passes in a no_grad context.
Could you explain why you are using these approaches and check if they are disallowing you from calling backward of the output?