diff --git a/interpolation/diffusion/timestep_sampler.py b/interpolation/diffusion/timestep_sampler.py
index a3f36984..072bdbab 100644
--- a/interpolation/diffusion/timestep_sampler.py
+++ b/interpolation/diffusion/timestep_sampler.py
@@ -125,7 +125,7 @@ def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
         self._loss_history = np.zeros(
             [diffusion.num_timesteps, history_per_term], dtype=np.float64
         )
-        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int64)
 
     def weights(self):
         if not self._warmed_up():
diff --git a/vsr/diffusion/timestep_sampler.py b/vsr/diffusion/timestep_sampler.py
index a3f36984..072bdbab 100644
--- a/vsr/diffusion/timestep_sampler.py
+++ b/vsr/diffusion/timestep_sampler.py
@@ -125,7 +125,7 @@ def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
         self._loss_history = np.zeros(
             [diffusion.num_timesteps, history_per_term], dtype=np.float64
         )
-        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int64)
 
     def weights(self):
         if not self._warmed_up():