class_docstr: Base class for RLlib's PyTorch models.
  
  This class defines the interface for RLlib's PyTorch models and checks
  whether inputs and outputs of forward are checked with `check_input_specs()` and
  `check_output_specs()` respectively.
  
  Example usage for a single Flattening layer:
  
  .. testcode::
  
      from ray.rllib.core.models.configs import ModelConfig
      from ray.rllib.core.models.torch.base import TorchModel
      import torch
  
      class FlattenModelConfig(ModelConfig):
          def build(self, framework: str):
              assert framework == "torch"
              return TorchFlattenModel(self)
  
      class TorchFlattenModel(TorchModel):
          def __init__(self, config):
              TorchModel.__init__(self, config)
              self.flatten_layer = torch.nn.Flatten()
  
          def _forward(self, inputs, **kwargs):
              return self.flatten_layer(inputs)
  
      model = FlattenModelConfig().build("torch")
      inputs = torch.Tensor([[[1, 2]]])
      print(model(inputs))
  
  .. testoutput::
  
      tensor([[1., 2.]])
class_docstr_tok: ["Base","class","for","RLlib","'s","PyTorch","models",".","This","class","defines","the","interface","for","RLlib","'s","PyTorch","models","and","checks","whether","inputs","and","outputs","of","forward","are","checked","with","`","check_input_specs","(",")","`","and","`","check_output_specs","(",")","`","respectively",".","Example","usage","for","a","single","Flattening","layer",":","..","testcode",":",":","from","ray.rllib.core.models.configs","import","ModelConfig","from","ray.rllib.core.models.torch.base","import","TorchModel","import","torch","class","FlattenModelConfig","(","ModelConfig",")",":","def","build","(","self",",","framework",":","str",")",":","assert","framework","==","``","torch","''","return","TorchFlattenModel","(","self",")","class","TorchFlattenModel","(","TorchModel",")",":","def","__init__","(","self",",","config",")",":","TorchModel.__init__","(","self",",","config",")","self.flatten_layer","=","torch.nn.Flatten","(",")","def","_forward","(","self",",","inputs",",","*","*","kwargs",")",":","return","self.flatten_layer","(","inputs",")","model","=","FlattenModelConfig","(",")",".build","(","``","torch","''",")","inputs","=","torch.Tensor","(","[","[","[","1",",","2","]","]","]",")","print","(","model","(","inputs",")",")","..","testoutput",":",":","tensor","(","[","[","1.",",","2",".","]","]",")"]
docstr: Initialized a TorchModel.
  
  Args:
      config: The ModelConfig to use.
docstr_tok: ["Initialized","a","TorchModel",".","Args",":","config",":","The","ModelConfig","to","use","."]
code: def __init__(self, config: ModelConfig):
      """Initialized a TorchModel.
  
      Args:
          config: The ModelConfig to use.
      """
      nn.Module.__init__(self)
      Model.__init__(self, config)
  
      # Raise errors if forward method is not decorated to check input specs.
      if not is_input_decorated(self.forward):
          raise ValueError(
              f"`{type(self).__name__}.forward()` not decorated with input "
              f"specification. Decorate it with @check_input_specs() to define a "
              f"specification and resolve this Error. If you don't want to check "
              f"anything, you can use an empty spec."
          )
  
      if is_output_decorated(self.forward):
          if log_once("torch_model_forward_output_decorated"):
              logger.warning(
                  f"`{type(self).__name__}.forward()` decorated with output "
                  f"specification. This is not recommended for torch models "
                  f"that are used with torch.compile() because it breaks "
                  f"torch dynamo's graph. This can lead lead to slower execution."
                  f"Remove @check_output_specs() from the forward() method to "
                  f"resolve this."
              )
  
file_path: ray/rllib/core/models/torch/base.py
imports: import abc
  import logging
  from typing import Tuple, Union
  import numpy
  from ray.rllib.core.models.base import Model
  from ray.rllib.core.models.configs import ModelConfig
  from ray.rllib.core.models.specs.checker import is_input_decorated, is_output_decorated, check_input_specs, check_output_specs
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.typing import TensorType
  from ray.util import log_once
label_desc: Use image node_id 1 to create a new TorchModel object from inherited base classes: nn, Model, abc with example: obj = TorchModel(config)