1from utils import NUM_LOOP_ITERS 2 3import torch 4 5 6def add_tensors_loop(x, y): 7 z = torch.add(x, y) 8 for i in range(NUM_LOOP_ITERS): 9 z = torch.add(z, x) 10 return z 11 12 13class SimpleAddModule(torch.nn.Module): 14 def __init__(self, add_op): 15 super().__init__() 16 self.add_op = add_op 17 18 def forward(self, x, y): 19 return self.add_op(x, y) 20