pytorch/issues/2230

bug.py

import torch
from torch.autograd import Variable
print('before import tf:')
print(torch.cuda.current_device(), torch.cuda.device_count())
import tensorflow as tf
with tf.device('/cpu:0'):
    emb = tf.Variable([[1,2],[3,4]], name="embedding")
print('after import tf:')
print(torch.cuda.current_device(), torch.cuda.device_count())
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
    sess.run(emb.initializer)
    print('after init:')
    print(torch.cuda.current_device(), torch.cuda.device_count())

print('after init (outside session):')
print(torch.cuda.current_device(), torch.cuda.device_count())

model = torch.nn.Linear(128, 1).cuda()
model = torch.nn.DataParallel(model).cuda()

data = Variable(torch.Tensor(8,128)).cuda()
x = model(data)