debug: 将归一化放在GPU上

This commit is contained in:
gouhanke
2026-02-11 17:13:55 +08:00
parent 320369ffb8
commit b42c1c68fd

View File

@@ -86,9 +86,8 @@ class SimpleRobotDataset(Dataset):
h5_path = f'observations/images/{cam_name}' h5_path = f'observations/images/{cam_name}'
if h5_path in f: if h5_path in f:
img = f[h5_path][meta["frame_idx"]] img = f[h5_path][meta["frame_idx"]]
img = torch.from_numpy(img).float() img = torch.from_numpy(img)
# 归一化到 [0, 1] 范围(与推理时保持一致) # 保持 uint8 格式以节省传输带宽,归一化移至 GPU (在 train_vla.py 中处理)
img = img / 255.0
frame[f"observation.{cam_name}"] = img.permute(2, 0, 1) # HWC -> CHW frame[f"observation.{cam_name}"] = img.permute(2, 0, 1) # HWC -> CHW
return frame return frame