test.py

import torch
import ray; ray.init()


@ray.remote(num_gpus=1) 
def foo(): 
  a = torch.rand(100000, device='cuda') 
  b = torch.rand(100000, device='cuda') 
  c = a * b * a * b * a * b * a * b  
  return c

results = [] 
for _ in range(10000): 
  results.append(foo.remote())
x = ray.get(results)