Hi Siddharth!



 SiddharthSingi:

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return x

Here isnt the layer fc2 a non leaf.


On the contrary, fc2.weight (and fc2.bias) are leaf tensors.
In the forward pass the values of fc2.weight do not depend on
the values of x (nor on the values in fc1).  They only change (in
the conventional use case) when you call optimizer.step().
Are you imagining that only the first layer is a leaf, that is that in:
    def forward(self, x):
        # x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return x

fc2 becomes a leaf because we commented out fc1?
To reiterate, the values of x of course depend upon whether fc1
was called, but the values of fc2.weight do not.
Best.
K. Frank