two_gpus.py

import numpy

import os

os.environ["CUDA_VISIBLE_DEVICES"]="0"
import torch

os.environ["CUDA_VISIBLE_DEVICES"]="1"
import tensorflow as tf

def main():

    for ii in range(4):
        print("PyTorch: Num GPUs Available: ", torch.cuda.device_count())

        ''' use pytorch on gpu0 '''
        with torch.cuda.device("cuda:0"):
            a = torch.from_numpy(numpy.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])).cuda()
            b = torch.from_numpy(numpy.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])).cuda()
            c = torch.matmul(a, b)
            print(c.device, c)

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

        ''' use tensorflow on gpu1 '''
        with tf.device("/device:GPU:0"):
            a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
            b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
            c = tf.matmul(a, b)
            print(c.device, c)



if __name__ == "__main__":
    main()