Ah! That’s it! Thank you for the solution.
I also create a small sample code to play around with this. As you can see, the lambda printouts the correct index only if I am using a local variable.
A weird thing is that when I try to pre-allocate a string message and pass it into the lambda function, the output is still 5.0. I dug into the internet and did not find out when is the string be evaluated, but I hypothesize that it is only evaluated when the __str()__ function is called. Therefore, it is keeping a reference to the cond, which has been set to 5.0 in the end (i.e. the same reason as the lambda function)
import torch
import torch.nn as nn


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        in_features = 5
        self.networks = nn.Sequential(
            nn.Linear(in_features=in_features, out_features=100, bias=True),
            nn.Linear(in_features=100, out_features=2),
        )
        self.initial_weights()

    def forward(self, input):
        num_conditions = 5

        conditions = torch.arange(start=1, end=num_conditions + 1, dtype=torch.float)
        out_features = {}
        for i, cond in enumerate(conditions):
            feature = torch.cat((input, cond.unsqueeze(0)), dim=0)
            feature = self.networks(feature)

            # message = "condition {}".format(cond)
            feature.register_hook(lambda grad, c=cond: print(c))
            out_features[i + 1] = feature
        return out_features

    def initial_weights(self):
        for module in self.networks:
            if isinstance(module, nn.Linear):
                if module.weight is not None:
                    nn.init.xavier_uniform_(module.weight, gain=1.0)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)


torch.manual_seed(42)
network = MyNet()
gt_cond = 3
gt = torch.ones(2)
data = torch.rand(4)
out_features = network(data)

loss = nn.functional.mse_loss(gt, out_features[gt_cond])
loss.backward()  # Expect to see 3.0 in the stdout