but you would have to train the last linear layer to give the correct outputs.

You mean the integer output layer that I want to add?

the range on the linear layers is infinite

Yeah, that’s the problem, it either goes to infinity or stays 0 if I apply sigmoid activation layer.
This is my current network:
x = self.activation(nn.Linear(16*16, 512)(x))
x = self.activation(nn.Linear(512, 256)(x))
x = self.activation(nn.Linear(256, 128)(x))
x = torch.log_softmax(nn.Linear(128, 65)(x).squeeze(), dim = 1)
# x = nn.Linear(65,1)(x)
 _, max_val = torch.max(x, dim = 1)
x = max_val.to(torch.float32)
# x = nn.Linear(1,1)(max_idx.to(torch.float32))

The last layer on which I’ve applied softmax, was my original last layer, I’ve 65 classes and it performs really well. But I am trying to convert that to integer output from the model, I tried a few things (commented lines) but none of them works. I just wanted to know if this can be done (and produce accurate results) at all or not.