install_driver_375.sh

# Installs 375.66 for CUDA 8 on Ubuntu 16.04

wget http://us.download.nvidia.com/tesla/375.66/nvidia-diag-driver-local-repo-ubuntu1604_375.66-1_amd64.deb
sudo dpkg -i nvidia-diag-driver-local-repo-ubuntu1604_375.66-1_amd64.deb
sudo apt-get update
sudo apt-get --allow-unauthenticated --assume-yes install cuda-drivers
sudo reboot now

install_driver_384.sh

# Installs 384.66 for CUDA 8 on Ubuntu 16.04

wget http://us.download.nvidia.com/tesla/384.66/nvidia-diag-driver-local-repo-ubuntu1604-384.66_1.0-1_amd64.deb
sudo dpkg -i nvidia-diag-driver-local-repo-ubuntu1604-384.66_1.0-1_amd64.deb
sudo apt-get update
sudo apt-get --allow-unauthenticated --assume-yes install cuda-drivers
sudo reboot now

leak.py

import resource
import torch
import torch.nn as nn
from torch.autograd import Variable

torch.backends.cudnn.benchmark = True

layers = [
  nn.Conv2d(3, 64, kernel_size=3, padding=1),
  nn.ReLU(),
  nn.Conv2d(64, 64, kernel_size=3, padding=1),
  nn.ReLU(),
  nn.Conv2d(64, 3, kernel_size=3, padding=1)
]
model = nn.Sequential(*layers).cuda()
loss_fn = nn.L1Loss().cuda()

num_iterations = 10000

N, C, H, W = 64, 3, 64, 64
for t in range(num_iterations):
  memory_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
  print('Iteration %d, using %.2f MB of memory'
        % (t, memory_kb / 1024))

  x = Variable(torch.randn(N, C, H, W).cuda())
  y = Variable(torch.randn(N, C, H, W).cuda())
  y_pred = model(x)
  loss = loss_fn(y_pred, y)
  loss.backward()
  model.zero_grad()

run.sh

# Install PyTorch 0.2 with pip
sudo apt-get --assume-yes install python3-venv
python3 -m venv env
source env/bin/activate
pip install http://download.pytorch.org/whl/cu80/torch-0.2.0.post1-cp35-cp35m-manylinux1_x86_64.whl 

# Run the test script
python leak.py