Hi,

I am training a ResNet18 with a graph neural network on top of it, I want to discard the GAT layer at evaluation, can someone help me with this?

```
def forward(self, x, lin=0, lout=5):
out = x
if lin < 1 and lout > -1:
out = self.conv1(out)
out = self.bn1(out)
out = F.relu(out)
if lin < 2 and lout > 0:
out = self.layer1(out)
if lin < 3 and lout > 1:
out = self.layer2(out)
if lin < 4 and lout > 2:
out = self.layer3(out)
if lin < 5 and lout > 3:
#This part
out = out.view(out.size(0), -1)
edges = self.topology(self.knn(out),512)
graph_data = (out,edges,edges)
out,edge,attention = self.gat_net(graph_data)
out = out.squeeze()
out = torch.reshape(out,(512,256,8,8))
# to here, I don't want it to perform at evaluation
out = self.layer4(out)
if lout > 4:
#out = F.avg_pool2d(out,2)
#res = out
out = F.adaptive_avg_pool2d(out, (1, 1))
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
```

Thanks