Thanks for the answer. Could you please make sure I understand it correctly? Is it correct to use it like this?
class Model(nn.Module):
    
    def __init__(self):
        super(Model, self).__init__()
        
        self.FeatureExtraction = ConvFeatureExtractor()
        self.rnn = RecurrentLayer()
        
    def forward(self, input): # input : [batch, step, channel, w, h]
        input = input.reshape(-1, 5, 3, 512, 512)
        
        prediction_outputs = []
        with torch.autograd.set_detect_anomaly(True):
          for i in range(5):
            image = input[:, i, :, :, :] # [batch, channel, w, h]
            visual_feature = self.FeatureExtraction(image)
            prediction = self.rnn(visual_feature)
            prediction_outputs.append(prediction)
        
        return prediction_outputs[-1]

model = Model()

x = torch.randn(10, 5, 3, 512, 512)
y = torch.randn(10, 1)

yhat = model(x)

mse = nn.MSELoss()
loss = mse(y, yhat)
loss.backward(retain_graph=True)

I’m running it on a Jupiter notebook, but the error message should pop up, but it doesn’t.