Torch.autograd.grad is returning `None` when calculating derivative wrt time

Abdul_Rehman: <p>I have an LSTM model that takes 3 sequences of temperature data and outputs the next sequence.</p>
<pre data-code-wrap="python"><code class="lang-python">input =&gt; [array([0.20408163, 0.40816327, 0.6122449 ]), 
          array([0.40816327, 0.6122449 , 0.81632653])]
 output=&gt; [tensor(0.81632653, dtype=torch.float64),
           tensor(0.91667510, dtype=torch.float64)]
</code></pre>
<p>Now,I want to combine this LSTM model with a Physics-Informed Neural Network (PINN) based on Newton’s Law of Cooling. The idea is to predict temperature using LSTM, then calculate the derivative of the predicted temperature with respect to time to incorporate the physics law into the loss function.</p>
<p>However, when I try to compute the gradient of the LSTM output with respect to time (t), the gradient returned is None. I’m not sure if I’m using torch.autograd correctly for this purpose.</p>
<p>Here is a simplified version of my code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

def create_lstm_model(input_size, hidden_size, num_layers, output_size):
    class LSTMModel(nn.Module):
        def __init__(self, input_size, hidden_size, num_layers, output_size):
            super(LSTMModel, self).__init__()
            self.hidden_size = hidden_size
            self.num_layers = num_layers
            self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
            self.fc = nn.Linear(hidden_size, output_size)

        def forward(self, x):
            h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
            c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
            out, _ = self.lstm(x, (h0, c0))
            out = self.fc(out[:, -1, :])
            return out


    return LSTMModel(input_size, hidden_size, num_layers, output_size)

def physics_loss_autograd(outputs, time_step):
    """
    Compute the physics-informed loss using autograd to get dT/dt.
    """
    # Compute dT/dt using autograd
    dT_dt = torch.autograd.grad(outputs, time_step, grad_outputs=torch.ones_like(outputs), create_graph=True)[0]

    # Newton's law of cooling: dT/dt = -k(T - T_ambient)
    residual = dT_dt + k * (outputs - T_ambient)

    # Physics loss is the L2 norm of the residual
    physics_loss = torch.mean(residual**2)

    return physics_loss


t = torch.arange(0,100) 
input_size =  1
hidden_size = 64
num_layers = 1
output_size = 1

# Create an instance of the LSTMModel using the function
model = create_lstm_model(input_size, hidden_size, num_layers, output_size)

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)


num_epochs = 100
for epoch in range(num_epochs):
    model.train()
    total_loss = 0

    for inputs, targets in train_loader:
        inputs, targets = inputs.float(), targets.float()  # Convert to float
        # print(inputs.shape)
        optimizer.zero_grad()
        outputs = model(inputs)


        
        data_loss = criterion(outputs, targets)
        
        # DO SOMETHING LIKE
        phys_loss = physics_loss_autograd(outputs, t)
        loss = data_loss + phys_loss

        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    if (epoch+1) % 20 == 0:
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(train_loader)}')
</code></pre>
<p>Has anyone worked on a similar problem? Any guidance on how to compute the temporal derivative of the LSTM output would be really helpful!</p>
<p>Additional Information:</p>
<ul>
<li>The time t is the timestep, not explicitly included as an input to the LSTM. I tried to include the timestep as additional feature to the input but its not calculating the derivative as well.</li>
<li>I’m still new to working with PINNs and combining physics loss with LSTM.</li>
</ul>
<p>The problem arises when I try to compute the physics-informed loss via torch.autograd.grad:</p>
<pre data-code-wrap="python"><code class="lang-python">dT_dt = torch.autograd.grad(outputs, t, grad_outputs=torch.ones_like(outputs), create_graph=True)[0]
</code></pre>
<p>This returns None for dT_dt. I suspect there’s an issue with how I’m handling the time_step or the autograd setup, but I’m not sure what exactly is going wrong.</p>

ptrblck: <aside class="quote no-group" data-username="Abdul_Rehman" data-post="1" data-topic="210960">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/abdul_rehman/48/73085_2.png" class="avatar"> Abdul_Rehman:</div>
<blockquote>
<p>The time t is the timestep, not explicitly included as an input to the LSTM.</p>
</blockquote>
</aside>
<p>If <code>t</code> is never used in the forward pass, it won’t have any effect on the output of the model and thus will also won’t have any valid gradients.</p>