Unable to export CFlow model to ONNX
Labels: module: onnx, triaged

### 🐛 Describe the bug

 
I am trying to export the [Cflow-ad](https://github.com/gudovskiy/cflow-ad) model to onnx format.
This model has an Encoder-Decoder architecture.
The Encoder is a feature extractor network (like Resnet18, MoblinetV3 etc..) that takes an image as an input, while the decoder  **loops** over the intermediate feature vectors of the encoder and does some specific "positional encoding" on the layers.
In the source code, the **decoder** is a list of decoders. 
In addition, there is no "Cflow model" class in the source code, therefore I wrote a simple Cflow module class to help export the model, as follows:

```
class CFlow(torch.nn.Module):

    def __init__(self, c):
        super(CFlow, self).__init__()
        L = c.pool_layers
        self.encoder, self.pool_layers, self.pool_dims = load_encoder_arch(c, L)
        self.encoder = self.encoder.to(c.device).eval()
        self.decoders = [load_decoder_arch(c, pool_dim) for pool_dim in self.pool_dims]
        self.decoders = [decoder.to(c.device) for decoder in self.decoders]
        params = list(self.decoders[0].parameters())
        for l in range(1, L):
            params += list(self.decoders[l].parameters())
        # optimizer
        self.optimizer = torch.optim.Adam(params, lr=c.lr)
        self.N=256

    def forward(self, x):
        P = c.condition_vec
        #print(self.decoders)
        self.decoders = [decoder.eval() for decoder in self.decoders]
        height = list()
        width = list()
        i=0
        test_dist = [list() for layer in self.pool_layers]
        test_loss = 0.0
        test_count = 0
        start = time.time()
        _ = self.encoder(x)
        with torch.no_grad():
            for l, layer in enumerate(self.pool_layers):
                e = activation[layer]  # BxCxHxW
                #
                B, C, H, W = e.size()
                S = H * W
                E = B * S
                #
                if i == 0:  # get stats
                    height.append(H)
                    width.append(W)
                #
                p = positionalencoding2d(P, H, W).to(c.device).unsqueeze(0).repeat(B, 1, 1, 1)
                c_r = p.reshape(B, P, S).transpose(1, 2).reshape(E, P)  # BHWxP
                e_r = e.reshape(B, C, S).transpose(1, 2).reshape(E, C)  # BHWxC

                decoder = self.decoders[l]
                FIB = E // self.N + int(E % self.N > 0)  # number of fiber batches
                for f in range(FIB):
                    if f < (FIB - 1):
                        idx = torch.arange(f * self.N, (f + 1) * self.N)
                    else:
                        idx = torch.arange(f * self.N, E)
                    #
                    c_p = c_r[idx]  # NxP
                    e_p = e_r[idx]  # NxC
                    # m_p = m_r[idx] > 0.5  # Nx1
                    #
                    if 'cflow' in c.dec_arch:
                        z, log_jac_det = decoder(e_p, [c_p, ])
                    else:
                        z, log_jac_det = decoder(e_p)
                    #
                    decoder_log_prob = get_logp(C, z, log_jac_det)
                    log_prob = decoder_log_prob / C  # likelihood per dim
                    loss = -log_theta(log_prob)
                    test_loss += t2np(loss.sum())
                    test_count += len(loss)
                    test_dist[l] = test_dist[l] + log_prob.detach().cpu().tolist()
        return height, width, test_dist
```
The argument 'c' is a config variable used to obtain useful parameters.

After the class is defined, I set all the needed parameters as in [main](https://github.com/gudovskiy/cflow-ad/blob/master/main.py).
Then, I added the following code to start the export:
```

model = CFlow(c)
print("Created !!!")
PATH = 'weights/mvtec_mobilenet_v3_large_freia-cflow_pl3_cb8_inp256_run0_Model_2022-11-08-10:50:39.pt'
model=load_weights(model,PATH)
model.eval()
batch_size = 1
x = torch.randn(batch_size, 3, 256, 256).to(c.device)
out=model(x)
```
The load_weights method is as follows:
```
def load_weights(model, filename):

    path = os.path.join(filename)
    state = torch.load(path)
    model.encoder.load_state_dict(state['encoder_state_dict'], strict=False)

    decoders = [decoder.load_state_dict(state, strict=False) for decoder, state in
                zip(model.decoders, state['decoder_state_dict'])]
    print('Loading weights from {}'.format(filename))
    return model
```
Finally:

```
torch.onnx.export(
        model,
        x,
        custom-d.onnx",
        export_params=True,
        verbose=True,
        opset_version=11,
        input_names=["input"],
        output_names=["output"],
    )
```
The code runs successfully, however no model is exported and thousands of lines are printed on the screen, like:
```
 9.5554e-02 -2.4600e-02  7.8328e-02  1.0419e-02
 4.5502e-02  1.4899e-03 -7.2213e-02 -3.1233e-02
 1.0151e-01 -2.4369e-02  5.9177e-03  6.9262e-02
 7.9299e-02 -3.7462e-02  7.7214e-03 -4.8352e-02
 3.6798e-02  6.1700e-02  6.9149e-02  3.8857e-02
 1.8279e-02  2.8520e-02  1.0628e-03 -7.4530e-03
[ torch.cuda.FloatTensor{296,148} ]
```
**What am I doing wrong? How can I export the model correctly?**
Check the full code [here](https://gist.github.com/Tekno-H/e114d6d843b3582f82c4a2a3cea1d46e)
I have been stuck on this problem for a while and can figure out how to solve it, any help is appreciated.
Thanks in advance.


### Versions

```
PyTorch version: 1.12.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.5.119
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1050 Ti
Nvidia driver version: 495.29.05
cuDNN version: Probably one of the following:
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.3
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] numpy                     1.23.3                   pypi_0    pypi
[conda] torch                     1.12.1                   pypi_0    pypi
[conda] torchvision               0.13.1                   pypi_0    pypi
```