Torchscript

[dibetgonzalez]

AllInOneBlock is 100% compatible with TorchScript?

Module 'AllInOneBlock' has no attribute 'w_0' :
File "/usr/local/lib/python3.9/site-packages/FrEIA/modules/all_in_one_block.py", line 173
'''Computes a permutation matrix from the reflection vectors that are
learned internally as nn.Parameters.'''
w = self.w_0
~~~~~~~~ <--- HERE
for vk in self.vk_householder:
w = torch.mm(w, torch.eye(self.in_channels).to(w.device) - 2 * torch.ger(vk, vk) / torch.dot(vk, vk))
'AllInOneBlock._construct_householder_permutation' is being compiled since it was called from 'AllInOneBlock.forward'
File "/usr/local/lib/python3.9/site-packages/FrEIA/modules/all_in_one_block.py", line 230
'''See base class docstring'''
if self.householder:
self.w_perm = self._construct_householder_permutation()
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
if rev or self.reverse_pre_permute:
self.w_perm_inv = self.w_perm.transpose(0, 1).contiguous()