Trying to call autograd.grad on a model inside of a custom autograd Function, works when initialized but not when weights/biases are set
class BernoulliApproximator(nn.Module): def __init__(self, hidden_dim): super().__init__() self.linear1 = nn.Linear(2, hidden_dim) self.linear2 = nn.Linear(hidden_dim, hidden_dim) self.linear3 = nn.Linear(hidden_dim, hidden_dim) self.linear4 = nn.Linear(hidden_dim, 1) self.relu = nn.ReLU() def forward(self, x): out = self.relu(self.linear1(x)) out = self.relu(self.linear2(out)) out = self.relu(self.linear3(out)) out = torch.sigmoid(self.linear4(out)) return out model = torch.load(‘bernoullimodel9.pth’,map_location=device) class BernoulliSampleFunction(torch.autograd.Function): @staticmethod def forward(ctx, probabilities, random_numbers): result = torch.zeros_like(probabilities) […]