class YourModel(nn.Module):
    
    def __init__(self, in_features, out_features):
        super(YourModel, self).__init__()
        self.fc = nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.fc(x)

model = YourModel(in_features, out_features)
model.fc.train(False)