typos

pts/modules/distribution_output.py

@@ -242,9 +242,10 @@ class MultivariateNormalOutput(DistributionOutput):
 
 class RealNVPOutput(DistributionOutput):
     def __init__(self, input_size, cond_size, n_blocks=3, n_hidden=1, hidden_size=100):
-        self.args_dim = {"cond": cond_label_size}
+        self.args_dim = {"cond": cond_size}
         self.dim = input_size
-        self.flow = RealNVP(nblocks=nblocks,
+        self.flow = RealNVP(
+            n_blocks=n_blocks,
             input_size=input_size,
             hidden_size=hidden_size,
             n_hidden=n_hidden,
@@ -252,7 +253,7 @@ class RealNVPOutput(DistributionOutput):
         )
     
     def domain_map(self, cond):
-        return cond
+        return (cond,)
 
     def distribution(self, distr_args, scale=None):
         cond, = distr_args
@@ -262,4 +263,4 @@ class RealNVPOutput(DistributionOutput):
 
     @property
     def event_shape(self) -> Tuple:
-        return (self.dim,)
+        return (self.dim,)

pts/modules/flows.py

@@ -176,12 +176,11 @@ class RealNVP(nn.Module):
         return self.net.inverse(u, self.cond)
 
     def log_prob(self, x):
-        u, sum_log_abs_det_jacobians = self.forward(x, self.cond)
+        u, sum_log_abs_det_jacobians = self.forward(x)
         return torch.sum(self.base_dist.log_prob(u) + sum_log_abs_det_jacobians, dim=-1)
 
     def sample(sample_shape=None):
-
-        u = self.base_dist.sample(sample_shape)
-        sample, _ = self.net.inverse(u, self.cond)
+        u = self.base_dist.sample(self.cond.shape)
+        sample, _ = self.inverse(u, self.cond)
 
         return sample