Computational graph with gradients?

[oseledets]

Thanks for the library!
Now the code takes an nn.Module and visualizes the forward pass;
Inside, it explicitly uses torch.no_grad.
However, if the forward pass of a module has autograd.grad inside, the library will fail.
Is it possible to modify the the library to allow for such use cases?

Simple example is attached.

import torch
import torch.nn as nn
# Define the model architecture
class visGrad(nn.Module):
    
    def __init__(self, model):
        super().__init__() # Add this line
        self.backbone =  model
        
    def forward(self, x, y):
        pred = self.backbone(x)
        er = pred - y
        loss = torch.sum(er**2)
        grd = torch.autograd.grad(loss, self.backbone.parameters())
        return loss

    
input_size = 1
hidden_size = 10
output_size = 1
model = nn.Sequential(
    nn.Linear(input_size, hidden_size),
    nn.ReLU(),
    nn.Linear(hidden_size, output_size),)
model.to(device)

    
x = torch.linspace(-2, 2, 256)
x = x[:, None]
y = x**2


model1 = visGrad(model)
model_graph = draw_graph(model1, input_data = (x, y), device='cpu')