Newer
Older
notebooks / python / test_gpu.py
#%%
%lsmagic

#%%
import torch
t_cpu = torch.rand(500,500,500)
%timeit t_cpu @ t_cpu

t_gpu = torch.rand(500,500,500).cuda()
%timeit t_gpu @ t_gpu

#TODO add TensorFlow GPU check