compare the GPU and CPU in matrix multiplication

GPU_vs_CPU.py

import torch
import time

def measure_time(device):
    # create random matrix
    size=1000 # size of matrix
    mat1=torch.randn(size,size, device=device)
    mat2=torch.randn(size,size , device=device)

    # measure the time taken for matrix multiplication
    start_time = time.time()
    result = torch.mm(mat1,mat2)
    end_time = time.time()

    return end_time - start_time



# run on CPU
cpu_time = measure_time("cpu")
print(f"Time on CPU:{cpu_time:.6f} second")
# print(f"Time on CPU: {cpu_time:.6f} seconds")

      
# check if GPU is available 
if torch.cuda.is_available():
    #Run on GPU
    gpu_time= measure_time("cuda")
    print(f"Time on GPU:{gpu_time:.6f}second")
else:
    print("GPU not available")