Simple add model — the add op will be delegated with CUDA device annotation.
| 76 | |
| 77 | |
| 78 | class ModuleAddWithDevice(nn.Module): |
| 79 | """Simple add model — the add op will be delegated with CUDA device annotation.""" |
| 80 | |
| 81 | def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: |
| 82 | return torch.add(a, b) |
| 83 | |
| 84 | def get_random_inputs(self): |
| 85 | return (torch.randn(2, 2), torch.randn(2, 2)) |
| 86 | |
| 87 | |
| 88 | def main() -> None: |