title: Use of the GPU on Colab with pytorch
body: <p>I am trying to use the GPU on Colab, but it takes about 400 seconds to train my model when I use GPU and 200 seconds to train it when I select CPU, so I think there is a problem somewhere..</p>
  <p>Where did I made a mistake ?</p>
  <pre><code>import torch
  from torch import Tensor
  import dlc_practical_prologue as prologue
  import time
  import matplotlib.pyplot as plt
  
  
  dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  print(dev)
  
  #%% Exercise1
  
  def sigma(x):
      return torch.tanh(x)
  
  def dsigma(x):
      return 1-torch.pow(torch.tanh(x), 2)
  
  #%% Exercise2
  
  def loss(v, t):    
      return torch.sum(torch.pow(t-v, 2))
  
  def dloss(v, t):
      #print(v.size(), t.size())
      return -2*(t-v)
  
  
  #%% Exercise3
  
  def forward_pass(w1, b1, w2, b2, x):
      s1 = w1.mv(x) + b1
      x1 = sigma(s1)
      #print(b1.size())
      s2 = w2.mv(x1) + b2
      x2 = sigma(s2)
      #print(s2.size(), x2.size())
      return x, s1 , x1, s2, x2
  
  
  def backward_pass(w1, b1, w2, b2, t, x, s1, x1, s2, x2, dl_dw1, dl_db1, dl_dw2, dl_db2):
      #print(dl_db2.size(), (dsigma(s2) * dloss(x2, t)).size())
      dl_db2_temp = dsigma(s2) * dloss(x2, t)
      dl_dw2_temp = dl_db2_temp.unsqueeze(0).t() * x1.unsqueeze(0)
      dl_db1_temp = w2.t().mv(dl_db2_temp)
      dl_dw1_temp = dl_db1_temp.unsqueeze(0).t() * x.unsqueeze(0)
      return dl_dw1 + dl_dw1_temp, dl_db1 + dl_db1_temp, dl_dw2 + dl_dw2_temp, dl_db2 + dl_db2_temp
  
  
  #%% Exercise4
  
  train_data, train_target, test_data, test_target = prologue.load_data(cifar = False, one_hot_labels = True, normalize = True, flatten = True)
  
  train_data = train_data.to(dev)
  train_target = train_target.to(dev)
  test_data = test_data.to(dev)
  test_target = test_target.to(dev)
  
  train_target = train_target * 0.9
  test_target = test_target * 0.9
  print(train_target.is_cuda)
  
  epsilon = 1e-6
  w1 = torch.empty((50, 784), device=dev).normal_(0, epsilon)
  b1 = torch.empty((50), device=dev).normal_(0, epsilon)
  w2 = torch.empty((10, 50), device=dev).normal_(0, epsilon)
  b2 = torch.empty((10), device=dev).normal_(0, epsilon)
  print(w1.is_cuda)
  
  dl_dw1 = torch.zeros((50, 784), device=dev)
  dl_db1 = torch.zeros((50), device=dev)  
  dl_dw2 = torch.zeros((10, 50), device=dev)  
  dl_db2 = torch.zeros((10), device=dev)  
  
  eta = 0.1 / train_target.size(0)
  nb_simul = 1000
  error_train = torch.zeros(nb_simul, device=dev)
  error_test = torch.zeros(nb_simul, device=dev)
  
  t1 = time.time()
  
  for i in range(nb_simul):
      dl_dw1 = torch.zeros((50, 784), device=dev)
      dl_db1 = torch.zeros((50), device=dev)  
      dl_dw2 = torch.zeros((10, 50), device=dev)  
      dl_db2 = torch.zeros((10), device=dev)
  
      for j in range(train_target.size(0)):
          x, s1 , x1, s2, x2 = forward_pass(w1, b1, w2, b2, train_data[j])
          #print(x.is_cuda)
          #print(x2.is_cuda)
          dl_dw1, dl_db1, dl_dw2, dl_db2 = backward_pass(w1, b1, w2, b2, train_target[j], x, s1, x1, s2, x2, dl_dw1, dl_db1, dl_dw2, dl_db2)
          #print(dl_dw1.is_cuda)
  
      w1 = w1 - eta * dl_dw1
      w2 = w2 - eta * dl_dw2
      b1 = b1 - eta * dl_db1
      b2 = b2 - eta * dl_db2
      #print(w1.is_cuda)
      
  
  
      for j in range(train_target.size(0)):
          x_train, s1_train, x1_train, s2_train, x2_train = forward_pass(w1, b1, w2, b2, train_data[j])
          _, predict = x2_train.max(0)
          _, value = train_target[j].max(0)
          if predict != value:
              error_train[i] += 1
      
      #print('error on training: ', error_train)
  
  
      for j in range(test_target.size(0)):
          x_test, s1_test, x1_test, s2_test, x2_test = forward_pass(w1, b1, w2, b2, test_data[j])
          _, predict = x2_test.max(0)
          _, value = test_target[j].max(0)
          if predict != value:
              error_test[i] += 1
  
      #print('error on test: ', error_test)
  
  
  t2 = time.time()
  print(t2-t1)
  
  #%% Plot
  
  abs_x = torch.arange(1, nb_simul+1, 1)
  plt.plot(abs_x, error_train)
  plt.plot(abs_x, error_test)
  print(error_train[-1])
  print(error_test[-1])
  </code></pre>