Create GPU load with PyTorch

pytorch-gpu-load.py

import multiprocessing
import torch
from torchvision.models import efficientnet_b0
from pytorch_benchmark import benchmark

def gpu(args):
    print(f"Starting workload on GPU{args}")
    model = efficientnet_b0().to(f"cuda:{args}")  # Model device sets benchmarking device
    sample = torch.randn(8, 3, 224, 224)  # (B, C, H, W)
    results = benchmark(model, sample, num_runs=1000)
    
if __name__ == '__main__':
    for i in range(6):
        p = multiprocessing.Process(target=gpu, args=(i,))
        p.start()