CUDA test script, verify there is working GPUs detected, by loading CUDA and training a tiny model

cuda-test.py

print("\nChecking CUDA with pyTorch\n")
import torch
import torch.nn as nn
dev = torch.device("cuda") if torch.cuda.is_available() else False 
if not dev:
    print("CUDA is not availale")
    raise SystemExit(1)

t1 = torch.randn(1,2)
t2 = torch.randn(1,2).to(dev)
t1.to(dev) 
t1 = t1.to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]], device='cuda:0') 
print("Tensor1 is CUDA: " + str(t1.is_cuda)) # True

class M(nn.Module):
    def __init__(self):        
        super().__init__()        
        self.l1 = nn.Linear(1,2)

    def forward(self, x):                      
        x = self.l1(x)
        return x
model = M()   # not on cuda
model.to(dev) # is on cuda (all parameters)
print("Model runs on CUDA: " + str(next(model.parameters()).is_cuda)) # True


print("IS CUDA Available: " + str(torch.cuda.is_available()))
print("Current CUDA Device: " + str(torch.cuda.current_device()))
print("First CUDA Devide: " + str(torch.cuda.device(0)))
print("CUDA Device count: " + str(torch.cuda.device_count()))
print("CUDA Device Name: " + str(torch.cuda.get_device_name(0)))


print("\nChecking CUDA with TensorFlow\n")

import tensorflow as tf

tf.config.list_physical_devices('GPU')
#with tf.device('/GPU:0'):
#      neural_network_1 = initialize_network_1()
#with tf.device('/GPU:1'):
#     neural_network_2 = initialize_network_2()


c1 = []
n = 10

def matpow(M, n):
    if n < 1: #Abstract cases where n < 1
        return M
    else:
        return tf.matmul(M, matpow(M, n-1))

with tf.device('/gpu:0'):
    a = tf.Variable(tf.random.uniform(shape=(10000, 10000)), name="a")
    b = tf.Variable(tf.random.uniform(shape=(10000, 10000)), name="b")
    c1.append(matpow(a, n))
    c1.append(matpow(b, n))

print("\nNum GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

TODO
Small test script - need to fix to give out meaningful output to machines - so instead/with printing - also use exit codes...

pip3.8 install torch==1.10.0