PyTorch训练Wav2vec2报错nll_loss_forward_reduce_cuda_kernel_2d_index未实现

错误原因 #

这个错误源于标签数据类型不匹配：PyTorch的NLLLoss（CrossEntropyLoss底层依赖）要求标签为torch.long（int64）类型，但Windows环境下加载数据集时，标签默认是Int32等非目标整数类型，和WSL的默认处理逻辑不一致，导致CUDA内核无法处理该类型。

具体解决步骤 #

• 强制转换标签数据类型
  在数据预处理函数中，显式将标签转为torch.long类型。修改你的预处理代码：

  def preprocess_function(examples):
      audio_arrays = [x["array"] for x in examples["audio"]]
      inputs = processor(audio_arrays, sampling_rate=16000, return_tensors="pt", padding=True)
      # 关键：将标签转为torch.long类型
      inputs["labels"] = torch.tensor(examples["label"], dtype=torch.long)
      return inputs
  

• 检查数据集标签列类型
  加载数据集后，先确认标签列的默认类型：

  from datasets import load_dataset
  dataset = load_dataset("你的数据集名称")
  print(dataset["train"].features["label"])
  

  如果输出是Value(dtype='int32')或其他非int64类型，除了预处理时转换，也可以在加载时指定类型：

  dataset = load_dataset("你的数据集名称", dtype={"label": "int64"})
  

• 验证训练数据的标签类型
  训练前手动检查一批数据的标签类型，确保符合要求：

  batch = next(iter(train_dataloader))
  print(batch["labels"].dtype)  # 预期输出：torch.int64
  

• 升级依赖版本（可选）
  Windows环境下部分旧版本的datasets和transformers存在跨平台类型处理差异，尝试升级到兼容版本：

  pip install --upgrade datasets>=2.10.0 transformers>=4.27.0
  

内容的提问来源于stack exchange，提问作者masimisa