Upload infer_gpu.py with huggingface_hub

infer_gpu.py

@@ -0,0 +1,39 @@
+import torch 
+import time
+import os
+import argparse
+import shutil
+import sys
+ 
+def parse_args():
+    parser = argparse.ArgumentParser(description='Matrix multiplication')
+    parser.add_argument('--gpus', help='gpu amount', required=True, type=int)
+    parser.add_argument('--size', help='matrix size', required=True, type=int)
+    parser.add_argument('--interval', help='sleep interval', required=True, type=float)
+    args = parser.parse_args()
+    return args
+ 
+ 
+def matrix_multiplication(args):
+ 
+    a_list, b_list, result = [], [], []    
+    size = (args.size, args.size)
+    
+    # for i in range(args.gpus):
+    
+    i = args.gpus
+    a_list.append(torch.rand(size, device=i))
+    b_list.append(torch.rand(size, device=i))
+    result.append(torch.rand(size, device=i))
+ 
+    while True:
+        # for i in range(args.gpus):
+
+        result[0] = a_list[0] * b_list[0]
+        time.sleep(args.interval)
+ 
+if __name__ == "__main__":
+    # usage: python matrix_multiplication_gpus.py --size 20000 --gpus 2 --interval 0.01
+    args = parse_args()
+    matrix_multiplication(args)
+