# -*- coding: utf-8 -*-

import os
import time
import torch
import argparse

parser = argparse.ArgumentParser(description='inf')
parser.add_argument('--gpu', default='1', type=str, help='index of gpu to use')
args = parser.parse_args()

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

n = 1000

x = torch.zeros(4, n, n).cuda()
rest_time = 0.0000000000001
while True:
    y = x * x
    time.sleep(rest_time)
    y1 = x * x