Запуск графического процессора с несколькими тензорными операторами PyTorch

Question

У меня следующий тензор PyTorch:

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6]])
X = torch.FloatTensor(X).cuda()

Мне было интересно, есть ли разница (особенно в скорости) между сценарием A или B ниже, если я запускаю следующие несколько операторов PyTorch в одной строке ?

Сценарий A:

X_sq_sum = (X**2).cuda().sum(dim = 1).cuda()

Сценарий B:

X_sq_sum = (X**2).sum(dim = 1).cuda()

ie. В сценарии A есть два .cuda(), тогда как в сценарии B только один .cuda().

Заранее большое спасибо.

Answer 1

Они будут работать одинаково, поскольку преобразование CUDA выполняется только один раз.

Как описано в docs , повторные вызовы .cuda() не будут выполняться, если объект уже находится в Память CUDA.