From 5c43fa1cceddc7cc58eb2e63698c5a36ea8a2eb5 Mon Sep 17 00:00:00 2001 From: Defalt Date: Sun, 19 Apr 2026 17:25:08 +0300 Subject: [PATCH] fix(policies): replace deprecated torch.cuda.amp.autocast with torch.amp.autocast (#3167) Co-authored-by: Steven Palma --- src/lerobot/policies/vqbet/vqbet_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/vqbet/vqbet_utils.py b/src/lerobot/policies/vqbet/vqbet_utils.py index f8bfcb06a..10180e0a4 100644 --- a/src/lerobot/policies/vqbet/vqbet_utils.py +++ b/src/lerobot/policies/vqbet/vqbet_utils.py @@ -27,7 +27,7 @@ import torch.distributed as distributed import torch.nn.functional as F # noqa: N812 from einops import pack, rearrange, reduce, repeat, unpack from torch import einsum, nn -from torch.cuda.amp import autocast +from torch.amp import autocast from torch.optim import Optimizer from .configuration_vqbet import VQBeTConfig @@ -1370,7 +1370,7 @@ class EuclideanCodebook(nn.Module): batch_samples = rearrange(batch_samples, "h ... d -> h (...) d") self.replace(batch_samples, batch_mask=expired_codes) - @autocast(enabled=False) + @autocast("cuda", enabled=False) def forward(self, x, sample_codebook_temp=None, mask=None, freeze_codebook=False): needs_codebook_dim = x.ndim < 4 sample_codebook_temp = (