Перебрать Dataloader, загруженный на графический процессор (MPS). ⇐ Python

[Гость]

Я работаю с PyTorch над M2 Max.

Я пытаюсь сократить время вычислений с помощью графического процессора.

У меня есть рабочий вариант с GPU:

mnist_test_loader = DataLoader(mnist_test_dataset, patch_size=32, shuffle=False) network.to(device="mps") для X в mnist_test_loader: X = X.to(device="mps") прогноз = сеть (X) Но дело в том, что при переборе Dataloader сначала X загружается на CPU и требуется время, чтобы переместить его на GPU X.to(device="mps")

Некоторое время я читал документацию и обнаружил, что можно загружать данные на графический процессор в Dataloader при его создании. Для этого я добавляю 2 параметра: pin_memory, pin_memory_device.

mnist_test_loader = DataLoader(mnist_test_dataset, patch_size=32, shuffle=False, pin_memory=Истина, pin_memory_device="mps") Идея состоит в том, что я передаю данные из Dataloader на графическом процессоре непосредственно в сеть, которая тоже находится на графическом процессоре, без необходимости перемещать их с процессора на графический процессор во время итерации.

Но тут возникает проблема:

для X вмокром_test_loader: ... Я не могу перебирать Dataloader на графическом процессоре. Когда я пытаюсь это сделать, я получаю следующую ошибку:

RuntimeError: попытка установить хранилище тензора на устройстве «cpu» в хранилище на другом устройстве «mps:0». Это больше не разрешено; устройства должны совпадать. Есть идеи, как с этим справиться? Может быть, есть другой способ получить данные из Dataloader? Спасибо, что уделили время.