fastai_onnx_wrapup_model.py

class ImageScale(nn.Module):
    def __init__(self): 
        super().__init__()
        self.denorminator = torch.full((3, sz, sz), 255.0, device=torch.device("cuda"))
    def forward(self, x): return torch.div(x, self.denorminator).unsqueeze(0)
    
# We need to:
# - Add ImageScale by 255.0 at the front
# - Replace LogSoftmax layer with Softmax at the end to get probability instead of loss/cost
final_model = [ImageScale()] + (list(learn.model.children())[:-1] + [nn.Softmax()])
final_model = nn.Sequential(*final_model)