diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py
index c9f98cf..5c1e11e 100644
--- a/pts/model/tempflow/tempflow_network.py
+++ b/pts/model/tempflow/tempflow_network.py
@@ -278,7 +278,7 @@ class TempFlowTrainingNetwork(nn.Module):
 
         return outputs, states, scale, lags_scaled, inputs
 
-    def distr(
+    def distr_args(
         self, rnn_outputs: torch.Tensor, scale: torch.Tensor,
     ):
         """
@@ -298,12 +298,13 @@ class TempFlowTrainingNetwork(nn.Module):
         distr_args
             Distribution arguments
         """
-        distr_args = self.proj_dist_args(rnn_outputs)
+        distr_args, = self.proj_dist_args(rnn_outputs)
 
-        # compute likelihood of target given the predicted parameters
-        distr = self.distr_output.distribution(distr_args, scale=scale)
+        # # compute likelihood of target given the predicted parameters
+        # distr = self.distr_output.distribution(distr_args, scale=scale)
 
-        return distr, distr_args
+        # return distr, distr_args
+        return distr_args
 
     def forward(
         self,
@@ -379,11 +380,11 @@ class TempFlowTrainingNetwork(nn.Module):
 
         # assert_shape(target, (-1, seq_len, self.target_dim))
 
-        distr, distr_args = self.distr(rnn_outputs=rnn_outputs, scale=scale)
+        distr_args = self.distr_args(rnn_outputs=rnn_outputs, scale=scale)
 
         # we sum the last axis to have the same shape for all likelihoods
         # (batch_size, subseq_length, 1)
-        likelihoods = -distr.log_prob(target).unsqueeze(-1)
+        likelihoods = -self.flow.log_prob(target, distr_args).unsqueeze(-1)
 
         # assert_shape(likelihoods, (-1, seq_len, 1))
 
@@ -412,7 +413,7 @@ class TempFlowTrainingNetwork(nn.Module):
 
         # self.distribution = distr
 
-        return (loss.mean(), likelihoods) + distr_args
+        return (loss.mean(), likelihoods, distr_args)
 
 
 class TempFlowPredictionNetwork(TempFlowTrainingNetwork):
@@ -494,12 +495,12 @@ class TempFlowPredictionNetwork(TempFlowTrainingNetwork):
                 unroll_length=1,
             )
 
-            distr, _ = self.distr(
+            distr_args = self.distr_args(
                 rnn_outputs=rnn_outputs, scale=repeated_scale,
             )
 
             # (batch_size, 1, target_dim)
-            new_samples = distr.sample()
+            new_samples = self.flow.sample(cond=distr_args)
 
             # (batch_size, seq_len, target_dim)
             future_samples.append(new_samples)
diff --git a/pts/modules/flows.py b/pts/modules/flows.py
index 01c8bfe..44c087c 100644
--- a/pts/modules/flows.py
+++ b/pts/modules/flows.py
@@ -142,7 +142,6 @@ class RealNVP(nn.Module):
         self.register_buffer('base_dist_mean', torch.zeros(input_size))
         self.register_buffer('base_dist_var', torch.ones(input_size))
 
-        self.__cond = None
         self.__scale = None
 
         # construct model
@@ -162,14 +161,6 @@ class RealNVP(nn.Module):
     def base_dist(self):
         return Normal(self.base_dist_mean, self.base_dist_var)
 
-    @property
-    def cond(self):
-        return self.__cond
-
-    @cond.setter
-    def cond(self, cond):
-        self.__cond = cond
-
     @property
     def scale(self):
         return self.__scale
@@ -178,27 +169,27 @@ class RealNVP(nn.Module):
     def scale(self, scale):
         self.__scale = scale
 
-    def forward(self, x):
+    def forward(self, x, cond):
         if self.scale is not None:
             x /= self.scale
-        return self.net(x, self.cond)
+        return self.net(x, cond)
 
-    def inverse(self, u):
-        x, log_abs_det_jacobian = self.net.inverse(u, self.cond)
+    def inverse(self, u, cond):
+        x, log_abs_det_jacobian = self.net.inverse(u, cond)
         if self.scale is not None:
             x *= self.scale
         return x, log_abs_det_jacobian
 
-    def log_prob(self, x):
-        u, sum_log_abs_det_jacobians = self.forward(x)
+    def log_prob(self, x, cond):
+        u, sum_log_abs_det_jacobians = self.forward(x, cond)
         return torch.sum(self.base_dist.log_prob(u) + sum_log_abs_det_jacobians, dim=-1)
 
-    def sample(self, sample_shape=torch.Size()):
-        if self.cond is not None:
-            shape = self.cond.shape[:-1] 
+    def sample(self, sample_shape=torch.Size(), cond=None):
+        if cond is not None:
+            shape = cond.shape[:-1] 
         else:
             shape = sample_shape
 
         u = self.base_dist.sample(shape)
-        sample, _ = self.inverse(u)
+        sample, _ = self.inverse(u, cond)
         return sample