pytorch/benchmarks/framework_overhead_benchmark/SimpleAddModule.py
Kimish Patel 2ba516d5b6 Added add op framework overhead benchmark for C2 (#23078)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23078

C2 benchmark.

Reviewed By: mingzhe09088

Differential Revision: D16122337

fbshipit-source-id: bf56e60c6e60eda2be2938d9f613708a4bc1669a
2019-07-22 11:27:50 -07:00

18 lines
472 B
Python

from __future__ import absolute_import, division, print_function, unicode_literals
import torch
from utils import NUM_LOOP_ITERS
def add_tensors_loop(x, y):
z = torch.add(x, y)
for i in range(NUM_LOOP_ITERS):
z = torch.add(z, x)
return z
class SimpleAddModule(torch.nn.Module):
def __init__(self, add_op):
super(SimpleAddModule, self).__init__()
self.add_op = add_op
def forward(self, x, y):
return self.add_op(x, y)