initializing apex allocates/reserves memory on wrong cuda-device

I try to initialize an Apex model on a non-default (not cuda:0) device but see allocation spiking up on other GPUs. When instantiating the model on cuda:1 the pytorch framework correctly allocates memory only on that device, however as soon as I execute amp.initialize the framework also allocates/reserves memory on cuda:0 . I attached a printout of a small repro and copied the code below.

import torch
import torch.nn as nn

import numpy as np

from apex import amp as amp

import pytorch_memlab as memlab
import pynvml as nvml

nvml.nvmlInit()
dev_handles= {
    0 : nvml.nvmlDeviceGetHandleByIndex(0),
    1 : nvml.nvmlDeviceGetHandleByIndex(1)}



rep = memlab.MemReporter()

rep.report()


dev = torch.device('cuda:1')



net = nn.Sequential(
    nn.Linear(1, 10),
    nn.ReLU(),
    nn.Linear(10, 1))
net.to(dev)

opt = torch.optim.SGD(net.parameters(), lr=1e-2)



rep.report()



for idx, h in dev_handles.items():
    mem_info = nvml.nvmlDeviceGetMemoryInfo(h)
    print(f'device cuda:{idx}, used/allocated memory: {mem_info.used / 1024**2} MB')


opt_level = 'O2'
net, opt = amp.initialize(net, opt, opt_level=opt_level, master_weights=False)


rep.report()


for idx, h in dev_handles.items():
    mem_info = nvml.nvmlDeviceGetMemoryInfo(h)
    print(idx, mem_info.used / 1024**2)



x = torch.randn(10, 1).to(dev)
y = net(x)


rep.report()


for idx, h in dev_handles.items():
    mem_info = nvml.nvmlDeviceGetMemoryInfo(h)
    print(idx, mem_info.used / 1024**2)

Start with the provided Python reproduction and inspect the amp.initialize entry point, focusing on device selection during initialization. Re-run the example while monitoring both CUDA devices; done means initializing a model on cuda:1 no longer allocates or reserves memory on cuda:0.