diff --git a/examples/longvit/utils.py b/examples/longvit/utils.py
index aa529ecf..2629cad6 100644
--- a/examples/longvit/utils.py
+++ b/examples/longvit/utils.py
@@ -21,7 +21,7 @@
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-from torch._six import inf
+from torch import inf
 # from torchmetrics import Metric
 from tensorboardX import SummaryWriter