Added permutation of AutoRegressiveNN to buffer
@@ -182,10 +182,11 @@ class ConditionalAutoRegressiveNN(nn.Module):
if permutation is None:
# By default set a random permutation of variables, which is important for performance with multiple steps
- self.permutation = torch.randperm(input_dim, device='cpu').to(torch.Tensor().device)
+ P = torch.randperm(input_dim, device='cpu').to(torch.Tensor().device)
else:
# The permutation is chosen by the user
- self.permutation = permutation.type(dtype=torch.int64)
+ P = permutation.type(dtype=torch.int64)
+ self.register_buffer('permutation', P)
# Create masks
self.masks, self.mask_skip = create_mask(
