diff --git a/util/misc.py b/util/misc.py
index 1d4e5eb13..ce4310dde 100644
--- a/util/misc.py
+++ b/util/misc.py
@@ -285,8 +285,7 @@ def __init__(self, tensors, mask: Optional[Tensor]):
         self.tensors = tensors
         self.mask = mask
 
-    def to(self, device):
-        # type: (Device) -> NestedTensor # noqa
+    def to(self, device: torch.device) -> 'NestedTensor':
         cast_tensor = self.tensors.to(device)
         mask = self.mask
         if mask is not None: