Compile models and losses with TorchScript for training and inference

[RMeli]

PR #11 did all the refactoring needed to compile the models with TorchScript but did not contain the last commit which implements the actual compilation. This PR pushes the missing commit.

[RMeli]

The branch

if self.reduction == "mean":
    reduced_loss = torch.mean(loss)
elif self.reduction == "sum":
    reduced_loss = torch.sum(loss)

results in the following PyTorch error:

E       RuntimeError: 
E       
E       reduced_loss is not defined in the false branch:
E         File "/biggin/b195/lina3015/Documents/git/gnina-torch/gnina/losses.py", line 91
E                   loss = diff * diff
E           
E               if self.reduction == "mean":
E               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
E                   reduced_loss = torch.mean(loss)
E                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
E               elif self.reduction == "sum":
E               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
E                   reduced_loss = torch.sum(loss)
E                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
E           
E               return self.scale * reduced_loss
E       and was used here:
E         File "/biggin/b195/lina3015/Documents/git/gnina-torch/gnina/losses.py", line 96
E                   reduced_loss = torch.sum(loss)
E           
E               return self.scale * reduced_loss
E                                   ~~~~~~~~~~~~ <--- HERE

The problem is that the if/elif block does not cover all execution paths leading to self.scale * reduced_loss and for the missing else clause the reduced_loss variable is undefined.

Using if/else is safe because of the assertion

gnina-torch/gnina/losses.py

Line 55 in 39985f3

assert reduction in ["mean", "sum"]

and solves the problem since reduced_loss is defined for all execution branches.

[codecov]

Codecov Report

Merging #19 (d34c1e6) into main (39985f3) will increase coverage by 0.00%.
The diff coverage is 100.00%.