##################################################
Enable GPU in Colab
1. Go to Runtime -> change runtime type (or "Runtime settings")
2. Under Hardware accelerator, select GPU (usually T4 GPU on the free tier)
3. click save
##################################################

import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

##################################################
import matplotlib.pyplot as plt
import numpy as np
import time

##################################################

## simlulate X
theseed = 34
n = 500000
p = 1000

np.random.seed(theseed)
X = np.random.normal(size=(n,p))

## compute X'X and time it
t1 = time.time()
%time xtx = (X.T) @ X
t2 = time.time()
print(f'time is {t2-t1}\n')
print(f'xtx00 is {xtx[0,0]}\n')

## make X a tensor on cpu
Xt = torch.from_numpy(X)

## compute X'X and time it
t1 = time.time()
%time xtx = (Xt.T) @ Xt
t2 = time.time()
print(f'time is {t2-t1}\n')
print(f'xtx00 is {xtx[0,0]}\n')

## put tensor on gpu
Xc = Xt.to('cuda')
## compute X'X and time it
t1 = time.time()
%time xtx = (Xc.T) @ Xc
t2 = time.time()
print(f'time is {t2-t1}\n')
print(f'xtx00 is {xtx[0,0]}\n')