diff --git a/PyCTBN/PyCTBN/structure_graph/trajectory_generator.py b/PyCTBN/PyCTBN/structure_graph/trajectory_generator.py
index 324a3ba..03aa634 100644
--- a/PyCTBN/PyCTBN/structure_graph/trajectory_generator.py
+++ b/PyCTBN/PyCTBN/structure_graph/trajectory_generator.py
@@ -55,7 +55,10 @@ class TrajectoryGenerator(object):
             t = time[next]
 
             if (max_tr != -1 and n_tr == max_tr) or (t_end != -1 and t >= t_end):
-                return Trajectory(self._importer.build_list_of_samples_array(sigma), len(self._vnames) + 1)
+                """ columns = self._importer.build_list_of_samples_array(sigma)
+                columns[0] = pd.to_numeric(columns[0])
+                return Trajectory(columns, len(self._vnames) + 1) """
+                return sigma
             else:
                 cim_row = np.array(cim[current_values.at[self._vnames[next]]])
                 cim_row[current_values.at[self._vnames[next]]] = 0
diff --git a/PyCTBN/PyCTBN/utility/abstract_importer.py b/PyCTBN/PyCTBN/utility/abstract_importer.py
index 8984b4d..0b03cd3 100644
--- a/PyCTBN/PyCTBN/utility/abstract_importer.py
+++ b/PyCTBN/PyCTBN/utility/abstract_importer.py
@@ -78,7 +78,7 @@ class AbstractImporter(ABC):
             Header of sample_frame = [Time | Variable values]
         """
         sample_frame = copy.deepcopy(sample_frame)
-        sample_frame.iloc[:, 0] = sample_frame.iloc[:, 0].diff().shift(-1)
+        #sample_frame.iloc[:, 0] = sample_frame.iloc[:, 0].diff().shift(-1)
         shifted_cols = sample_frame[columns_header].shift(-1).fillna(0).astype('int32')
         shifted_cols.columns = shifted_cols_header
         sample_frame = sample_frame.assign(**shifted_cols)
diff --git a/test_trajectorygenerator.py b/test_trajectorygenerator.py
index b0535c2..ad13b6b 100644
--- a/test_trajectorygenerator.py
+++ b/test_trajectorygenerator.py
@@ -1,6 +1,7 @@
 import unittest
 import random
 
+from PyCTBN.PyCTBN.structure_graph.trajectory import Trajectory
 from PyCTBN.PyCTBN.structure_graph.trajectory_generator import TrajectoryGenerator
 from PyCTBN.PyCTBN.utility.json_importer import JsonImporter
 
@@ -22,17 +23,19 @@ class TestTrajectoryGenerator(unittest.TestCase):
         tg = TrajectoryGenerator(self.j1)
         end_time = random.randint(5, 100)
         sigma = tg.CTBN_Sample(end_time)
-        self.assertLessEqual(sigma.times[len(sigma.times) - 1], end_time)
-        for index in range(len(sigma.times)):
+        traj = Trajectory(self.j1.build_list_of_samples_array(sigma), len(self.j1.sorter) + 1)
+        self.assertLessEqual(traj.times[len(traj.times) - 1], end_time)
+        for index in range(len(traj.times)):
             if index > 0:
-                self.assertLess(sigma.times[index - 1], sigma.times[index])
-                diff = abs(sum(sigma.trajectory[index - 1]) - sum(sigma.trajectory[index]))
+                self.assertLess(traj.times[index - 1], traj.times[index])
+                diff = abs(sum(traj.trajectory[index - 1]) - sum(traj.trajectory[index]))
                 self.assertEqual(diff, 1)
 
     def test_generated_trajectory_max_tr(self):
         tg = TrajectoryGenerator(self.j1)
         n_tr = random.randint(5, 100)
         sigma = tg.CTBN_Sample(max_tr = n_tr)
-        self.assertEqual(len(sigma.times), n_tr + 1)
+        traj = Trajectory(self.j1.build_list_of_samples_array(sigma), len(self.j1.sorter) + 1)
+        self.assertEqual(len(traj.times), n_tr + 1)
 
 unittest.main()
\ No newline at end of file