Maybe as simple as adding a flag to the class, self.mode=0 then
   def forward(self, x):
       identity = x
        if self.mode:
            with torch.no_grad():
                out = self.conv1(x)
                out = self.bn1(out)
                out = self.relu(out)
                out = self.conv2(out)
                out = self.bn2(out)
        else:
            out = self.conv1(x)
            out = self.bn1(out)
            out = self.relu(out)
            out = self.conv2(out)
            out = self.bn2(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        out = self.relu1(out)
        return out

the modified resnet does not fail; i havent checked what gradients are doing yet, but seems should work… if anyone sees any issue flag me. thanks.