| 524 | extra_link_args=self.strip_empty_entries(self.extra_ldflags())) |
| 525 | |
| 526 | def load(self, verbose=False): |
| 527 | if self.name in __class__._loaded_ops: |
| 528 | return __class__._loaded_ops[self.name] |
| 529 | |
| 530 | from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name |
| 531 | from deepspeed.accelerator import get_accelerator |
| 532 | if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: |
| 533 | # Ensure the op we're about to load was compiled with the same |
| 534 | # torch/cuda versions we are currently using at runtime. |
| 535 | self.validate_torch_version(torch_info) |
| 536 | if torch.cuda.is_available() and isinstance(self, CUDAOpBuilder): |
| 537 | self.validate_torch_op_version(torch_info) |
| 538 | |
| 539 | op_module = importlib.import_module(self.absolute_name()) |
| 540 | __class__._loaded_ops[self.name] = op_module |
| 541 | return op_module |
| 542 | else: |
| 543 | return self.jit_load(verbose) |
| 544 | |
| 545 | def jit_load(self, verbose=True): |
| 546 | if not self.is_compatible(verbose): |