diff --git a/pyFTS/models/pwfts.py b/pyFTS/models/pwfts.py
index 6e31a41..53da7f6 100644
--- a/pyFTS/models/pwfts.py
+++ b/pyFTS/models/pwfts.py
@@ -505,9 +505,9 @@ class ProbabilisticWeightedFTS(ifts.IntervalFTS):
 
             # Find all bins of past distributions with probability greater than zero
 
-            for ct, d in enumerate(self.lags):
-                dd = ret[k - d]
-                vals = [float(v) for v in dd.bins if round(dd.density(v), 4) > 0]
+            for ct, lag in enumerate(self.lags):
+                dd = ret[k - lag]
+                vals = [float(v) for v in dd.bins if np.round(dd.density(v), 4) > 0.0]
                 lags.append( sorted(vals) )
 
 
@@ -516,9 +516,8 @@ class ProbabilisticWeightedFTS(ifts.IntervalFTS):
             for path in product(*lags):
 
                 # get the combined probabilities for this path
-
-                pk = np.prod([ret[k - self.max_lag + o].density(path[ct])
-                              for ct, o in enumerate(self.lags)])
+                pk = np.prod([ret[k - (self.max_lag + lag)].density(path[ct])
+                              for ct, lag in enumerate(self.lags)])
 
 
                 d = self.forecast_distribution(path)[0]
diff --git a/pyFTS/tests/general.py b/pyFTS/tests/general.py
index 02ef077..d0531bd 100644
--- a/pyFTS/tests/general.py
+++ b/pyFTS/tests/general.py
@@ -32,7 +32,9 @@ pfts1_taiex.fit(dataset[:train_split], save_model=True, file_path='pwfts', order
 pfts1_taiex.shortname = "1st Order"
 #print(pfts1_taiex)
 
-tmp = pfts1_taiex.predict(dataset[train_split:train_split+200], type='distribution')
+#tmp = pfts1_taiex.predict(dataset[train_split:train_split+200], type='distribution')
+
+tmp = pfts1_taiex.predict(dataset[train_split:train_split+200], type='distribution', steps_ahead=20)
 
 '''
 #dataset = SP500.get_data()[11500:16000]