Usage
state_dict = torch.load('model.pth', map_location='cpu')
model.load_state_dict(state_dict, strict=False)
model.eval()
def predict(im):
x = torch.tensor(im, dtype=torch.float32)
with torch.no_grad():
out = model(x)
probabilities = F.softmax(out[0], dim=0)
