diff --git a/tape/training.py b/tape/training.py
index 2071485a..84ad3677 100644
--- a/tape/training.py
+++ b/tape/training.py
@@ -132,7 +132,7 @@ def __init__(self,
         self.max_grad_norm = max_grad_norm
         self._global_step = 0
         self._local_rank = local_rank
-        self._overflow_buf = torch.cuda.IntTensor([0])  # type: ignore
+        self._overflow_buf = torch.IntTensor([0])  # type: ignore
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self._delay_accumulation = fp16 and local_rank != -1