Issues using the torch.autograd.grad in the tinycuda environment

[joejep]

gradients = torch.autograd.grad( outputs=potential.sum(), inputs=flat_points, create_graph=create_graph, retain_graph=True, only_inputs=True, allow_unused=True )[0]

is there a way out. I am forced to use the traditional architecture(using MLP) self.block1 = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(),

    )

instead of using the

network_config = { "otype": "FullyFusedMLP", "activation": "ReLU", "output_activation": "Sigmoid", # Final activation for RGB "n_neurons": hidden_dim, "n_hidden_layers": 8 # Combined layers from all blocks }

    self.network = tcnn.Network(
       n_input_dims=total_input_dim,
        n_output_dims=3,  # RGB output
        network_config=network_config

I get an error using tinycuda when the torch.autograd.grad is trying to create graph but no error using the nn.sequential. I am not okay with inn.sequential because i need to speed up the training. Any help will be appreciated