investigate if conditions in torch script

[TomMelt]

Currently #33 hacks the utils/model_definition.py to disable an if condition that isn't used for 1x1 model but causes issue when trying to script

nonlocal_gwfluxes/utils/model_definition.py

Lines 74 to 78 in c389d75

	# if self.fac == 1:
	# x = torch.squeeze(self.dropout0(self.act_cnn(self.conv1(x))))
	# elif self.fac == 2:
	# x = torch.squeeze(self.dropout0(self.act_cnn(self.conv1(x))))
	# x = torch.squeeze(self.dropout0_2(self.act_cnn2(self.conv2(x))))

Error:

RuntimeError: 
Module 'ANN_CNN' has no attribute 'dropout0' :
  File "/home/melt/sync/cambridge/projects/current/nlgw-cam/nonlocal_gwfluxes/era5_training/../utils/model_definition.py", line 75
    def forward(self, x):
        if self.fac == 1:
            x = torch.squeeze(self.dropout0(self.act_cnn(self.conv1(x))))
                              ~~~~~~~~~~~~~ <--- HERE
        elif self.fac == 2:
            x = torch.squeeze(self.dropout0(self.act_cnn(self.conv1(x))))

• investigate cause of error
• find permanent solution rather than commenting the code out

[TomMelt]

According to the PyTorch docs we will have to modify the code:

TorchScript supports a subset of Python’s variable resolution (i.e. scoping) rules. Local variables behave the same as in Python, except for the restriction that a variable must have the same type along all paths through a function. If a variable has a different type on different branches of an if statement, it is an error to use it after the end of the if statement.
Similarly, a variable is not allowed to be used if it is only defined along some paths through the function.

For example, the following code:

@torch.jit.script
def foo(x):
    if x < 0:
        y = 4
    print(y)

will give this error

Traceback (most recent call last):
  ...
RuntimeError: ...

y is not defined in the false branch...
@torch.jit.script...
def foo(x):
    if x < 0:
    ~~~~~~~~~
        y = 4
        ~~~~~ <--- HERE
    print(y)
and was used here:
    if x < 0:
        y = 4
    print(y)
          ~ <--- HERE...

[TomMelt]

@TomMelt to look into proper work around by pre-defining unused vars