How can I adapt INN to model with only mlp layers?

[Jaehyung-Lim]

First, thanks for sharing your wonderful research.

I have a tiny question about this project during applying INN to MLP model.

I'm implementing MLP layer (64 -> 32 -> 16 -> 1).

Their names are 'fc1', 'fc2', 'fc3', namely there weights are 'fc1.weight', 'fc2.weight', 'fc3.weight'.

Then, How can I use INN wrapper with this MLP model?

I want you to kindly share example of "continuous_dims" and "example_input".

Thank you again!

[b1n0]

Hello! Thanks for your interest in our research. Sure, it is quite simple:
example_input = torch.randn(1, 64)
continuous_dims = { 'fc1.weight': [0, 1], 'fc2.weight': [0, 1], 'fc3.weight': [1], }
here 0 corresponds to the out_channels dimension of fc layer and 1 to the in_channels dim.