diff --git a/trajectron/environment/scene_graph.py b/trajectron/environment/scene_graph.py
index 1113bd4..4d969e9 100644
--- a/trajectron/environment/scene_graph.py
+++ b/trajectron/environment/scene_graph.py
@@ -135,10 +135,10 @@ class TemporalSceneGraph(object):
         position_cube = np.full((total_timesteps, N, 2), np.nan)
 
         adj_cube = np.zeros((total_timesteps, N, N), dtype=np.int8)
-        dist_cube = np.zeros((total_timesteps, N, N), dtype=np.float)
+        dist_cube = np.zeros((total_timesteps, N, N), dtype=float)
 
         node_type_mat = np.zeros((N, N), dtype=np.int8)
-        node_attention_mat = np.zeros((N, N), dtype=np.float)
+        node_attention_mat = np.zeros((N, N), dtype=float)
 
         for node_idx, node in enumerate(nodes):
             if online: