Home
last modified time | relevance | path

Searched refs:MyModuleInterface (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/nn/jit/
H A Dtest_instantiator.py22 class MyModuleInterface: class
43 ) = instantiator.get_arg_return_types_from_interface(MyModuleInterface)
62 MyModuleInterface
/aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/rpc/jit/
H A Drpc_test.py355 class MyModuleInterface(torch.nn.Module): class
380 return rpc.RRef(MyScriptModule(a), type_hint=MyModuleInterface)
389 def script_rref_run_forward_my_script_module(rref: RRef[MyModuleInterface]) -> Tensor: argument
410 rref_script_module = rpc.RRef(MyScriptModule(self.rank), MyModuleInterface)
466 def use_rref_on_owner(rref: RRef[MyModuleInterface]) -> Tensor: argument
820 def my_script_module_init(rank: int) -> MyModuleInterface:
825 def construct_my_script_module(rank: int) -> MyModuleInterface:
831 ref_script_module: RRef[MyModuleInterface], t: Tensor argument
/aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/nn/api/
H A Dremote_module_test.py55 class MyModuleInterface: class
123 _module_interface_cls=MyModuleInterface,
208 def run_forward(scripted_remote_module: MyModuleInterface): argument
726 def run_forward(scripted_remote_module: MyModuleInterface): argument