| import torch |
| import argparse |
| from common import SubTensor, WithTorchFunction, SubWithTorchFunction # noqa: F401 |
| |
| Tensor = torch.Tensor |
| |
| NUM_REPEATS = 1000000 |
| |
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser( |
| description="Run the torch.add for a given class a given number of times." |
| ) |
| parser.add_argument( |
| "tensor_class", metavar="TensorClass", type=str, help="The class to benchmark." |
| ) |
| parser.add_argument( |
| "--nreps", "-n", type=int, default=NUM_REPEATS, help="The number of repeats." |
| ) |
| args = parser.parse_args() |
| |
| TensorClass = globals()[args.tensor_class] |
| NUM_REPEATS = args.nreps |
| |
| t1 = TensorClass(1) |
| t2 = TensorClass(2) |
| |
| for _ in range(NUM_REPEATS): |
| torch.add(t1, t2) |
