FrEIA forward compatibility

[wltjr1007]

Since FrEIA went through few major updates, this repository do not work with the latest FrEIA.

Are there any quick fix for forward compatibility?

For the time being, I am using the version of FrEIA with framework.py in the root directory (i.e. https://github.com/VLL-HD/FrEIA/tree/550257b10af7d8772b08d4aa9b18772e2c0225f9)

[wltjr1007]

Use with latest FrEIA gives me the following error:

Traceback (most recent call last):
  File "/DataCommon/jsyoon/IB-INN/FrEIA/framework/graph_inn.py", line 298, in forward
    mod_out = node.module(mod_in, rev=rev, jac=jac)
  File "/home/jsyoon/anaconda3/envs/ib-inn/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'jac'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "main.py", line 29, in <module>
    train.train(args)
  File "/DataCommon/jsyoon/IB-INN/train.py", line 98, in train
    losses = inn(x, y)
  File "/home/jsyoon/anaconda3/envs/ib-inn/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/DataCommon/jsyoon/IB-INN/model.py", line 141, in forward
    z = self.inn(x)
  File "/home/jsyoon/anaconda3/envs/ib-inn/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/DataCommon/jsyoon/IB-INN/FrEIA/framework/reversible_graph_net.py", line 36, in forward
    return super().forward(x_or_z, c, rev, jac, intermediate_outputs)
  File "/DataCommon/jsyoon/IB-INN/FrEIA/framework/graph_inn.py", line 300, in forward
    raise RuntimeError(f"{node} encountered an error.") from e
RuntimeError: Node 'CONV_0_0': [(3, 32, 32)] -> AIO_Block -> [(3, 32, 32)] encountered an error.

[ardizzone]

Hi! Sorry for the late answer!
The code runs with the 0.2 FrEIA version (see README, I just updated it)
I will try to fix it to work with the current FrEIA by next week!