pytorch-gputest.py

import torch
import torch.nn as nn
from tqdm import tqdm
device = torch.device('cuda')

# 5629MiB / 11178MiB
# NVIDIA-SMI 410.78       Driver Version: 410.78       CUDA Version: 10.0

torch.backends.cudnn.benchmark = True

x = torch.randn(128, 1024, 8, 8).to(device)
fc = nn.Sequential(
	nn.Conv2d(1024, 1024, 3, padding=1),
	nn.ReLU(inplace=True),
	nn.Conv2d(1024, 1024, 3, padding=1),
	nn.ReLU(inplace=True),
	nn.Conv2d(1024, 1024, 3, padding=1),
	nn.ReLU(inplace=True),
).to(device)

for _ in tqdm(range(100000)):
	fc.zero_grad()
	y = fc(x)
	loss = y.mean()
	loss.backward()