from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

# pytorch 0.4.1

import pdb
import time
import os,sys

import torch
import torch.nn as nn

dtype = torch.float
device = torch.device("cuda")

class TEST(nn.Module):
  def __init__(self, b):
    super(TEST, self).__init__()
    self.b = b
  def forward(self, A):
    b_cuda = self.b.cuda(device=A.device)
    torch.cuda.synchronize()
    start = time.time()
    X, _ = torch.gesv(b_cuda, A)
    torch.cuda.synchronize()
    print(time.time()-start)
    return X


if __name__ == '__main__':


  os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"

  A = torch.rand(30, 1024, 1024, device=device, dtype=dtype)
  b = torch.rand(10, 1024, 1)

  model = TEST(b)
  model = nn.DataParallel(model).cuda()

  for _ in range(10):
    z = model(A)


  '''
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"

  A = torch.rand(10, 1024, 1024, device=device, dtype=dtype)
  b = torch.rand(10, 1024, 1)

  model = TEST(b)
  model = nn.DataParallel(model).cuda()

  for _ in range(10):
    z = model(A)

  ''"


Hi, here is my code. When I use multi-gpu, the output time consuming is about 200ms, however, when I use single GPU like below, the output time consuming is about 145ms, I do not know why? Test on Titan X.