diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py
index 646d2b6..b5953b5 100644
--- a/pts/modules/distribution_output.py
+++ b/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,)
\ No newline at end of file
+        return (self.dim,)
diff --git a/pts/modules/flows.py b/pts/modules/flows.py
index d2e893a..1337ec9 100644
--- a/pts/modules/flows.py
+++ b/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
\ No newline at end of file