class_docstr: A base class for RLlib PyTorch RLModules.
  
  Note that the `_forward` methods of this class can be 'torch.compiled' individually:
      - `TorchRLModule._forward_train()`
      - `TorchRLModule._forward_inference()`
      - `TorchRLModule._forward_exploration()`
  
  As a rule of thumb, they should only contain torch-native tensor manipulations,
  or otherwise they may yield wrong outputs. In particular, the creation of RLlib
  distributions inside these methods should be avoided when using `torch.compile`.
  When in doubt, you can use `torch.dynamo.explain()` to check whether a compiled
  method has broken up into multiple sub-graphs.
  
  Compiling these methods can bring speedups under certain conditions.
class_docstr_tok: ["A","base","class","for","RLlib","PyTorch","RLModules",".","Note","that","the","`","_forward","`","methods","of","this","class","can","be","'torch.compiled","'","individually",":","-","`","TorchRLModule._forward_train","(",")","`","-","`","TorchRLModule._forward_inference","(",")","`","-","`","TorchRLModule._forward_exploration","(",")","`","As","a","rule","of","thumb",",","they","should","only","contain","torch-native","tensor","manipulations",",","or","otherwise","they","may","yield","wrong","outputs",".","In","particular",",","the","creation","of","RLlib","distributions","inside","these","methods","should","be","avoided","when","using","`","torch.compile","`",".","When","in","doubt",",","you","can","use","`","torch.dynamo.explain","(",")","`","to","check","whether","a","compiled","method","has","broken","up","into","multiple","sub-graphs",".","Compiling","these","methods","can","bring","speedups","under","certain","conditions","."]
code: def load_state(self, dir: Union[str, pathlib.Path]) -> None:
      path = str(pathlib.Path(dir) / self._module_state_file_name())
      self.set_state(torch.load(path))
  
imports: import pathlib
  from typing import Any, List, Mapping, Tuple, Union, Type
  from packaging import version
  from ray.rllib.core.rl_module import RLModule
  from ray.rllib.core.rl_module.rl_module_with_target_networks_interface import RLModuleWithTargetNetworksInterface
  from ray.rllib.core.rl_module.torch.torch_compile_config import TorchCompileConfig
  from ray.rllib.models.torch.torch_distributions import TorchDistribution
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor, TORCH_COMPILE_REQUIRED_VERSION
  from ray.rllib.utils.typing import NetworkType
label_desc: Use image node_id 8 for calling the TorchRLModule obj's underlying member method code with example usage: obj.load_state(dir) without return types