Revert "Make Context to be Device-agnostic Step by Step (1/N) (#136519)"

This reverts commit 4a8e49389c33934234dc89616fd17a58e760e2e7.

Reverted https://github.com/pytorch/pytorch/pull/136519 on behalf of https://github.com/clee2000 due to breaking internal tests related to MITA, @ezyang has a forward fix? ([comment](https://github.com/pytorch/pytorch/pull/136519#issuecomment-2414588302))
This commit is contained in:
PyTorch MergeBot
2024-10-15 17:19:16 +00:00
parent 9af4e0d2aa
commit d4d687ffb2
31 changed files with 90 additions and 110 deletions

View File

@ -75,7 +75,7 @@ class NCCLTest : public NCCLTestBase {
int inputDim = 3)
: NCCLTestBase(path, pgTimeout), rank_(rank), worldSize_(worldSize) {
// Each device has a single tensor to perf the NCCL op
::at::globalContext().lazyInitDevice(c10::DeviceType::CUDA);
::at::globalContext().lazyInitCUDA();
tensors_.resize(numDevices_);
inputs_.resize(numDevices_);
outputs_.resize(numDevices_);