class_docstr: PyTorch specific Policy class to use with RLlib.
class_docstr_tok: ["PyTorch","specific","Policy","class","to","use","with","RLlib","."]
code: def load_batch_into_buffer(
      self,
      batch: SampleBatch,
      buffer_index: int = 0,
  ) -> int:
      # Set the is_training flag of the batch.
      batch.set_training(True)
  
      # Shortcut for 1 CPU only: Store batch in `self._loaded_batches`.
      if len(self.devices) == 1 and self.devices[0].type == "cpu":
          assert buffer_index == 0
          pad_batch_to_sequences_of_same_size(
              batch=batch,
              max_seq_len=self.max_seq_len,
              shuffle=False,
              batch_divisibility_req=self.batch_divisibility_req,
              view_requirements=self.view_requirements,
          )
          self._lazy_tensor_dict(batch)
          self._loaded_batches[0] = [batch]
          return len(batch)
  
      # Batch (len=28, seq-lens=[4, 7, 4, 10, 3]):
      # 0123 0123456 0123 0123456789ABC
  
      # 1) split into n per-GPU sub batches (n=2).
      # [0123 0123456] [012] [3 0123456789 ABC]
      # (len=14, 14 seq-lens=[4, 7, 3] [1, 10, 3])
      slices = batch.timeslices(num_slices=len(self.devices))
  
      # 2) zero-padding (max-seq-len=10).
      # - [0123000000 0123456000 0120000000]
      # - [3000000000 0123456789 ABC0000000]
      for slice in slices:
          pad_batch_to_sequences_of_same_size(
              batch=slice,
              max_seq_len=self.max_seq_len,
              shuffle=False,
              batch_divisibility_req=self.batch_divisibility_req,
              view_requirements=self.view_requirements,
          )
  
      # 3) Load splits into the given buffer (consisting of n GPUs).
      slices = [
          slice.to_device(self.devices[i])
          for i, slice in enumerate(slices)
      ]
      self._loaded_batches[buffer_index] = slices
  
      # Return loaded samples per-device.
      return len(slices[0])
  
imports: import copy
  import functools
  import logging
  import math
  import os
  import threading
  import time
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
  import gymnasium
  import numpy
  import tree
  import ray
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.policy import Policy, PolicyState
  from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils import NullContextManager, force_list
  from ray.rllib.utils.annotations import DeveloperAPI, override
  from ray.rllib.utils.error import ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics import DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, NUM_AGENT_STEPS_TRAINED, NUM_GRAD_UPDATES_LIFETIME
  from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.spaces.space_utils import normalize_action
  from ray.rllib.utils.threading import with_lock
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
  from ray.rllib.utils.typing import AlgorithmConfigDict, GradInfoDict, ModelGradients, ModelWeights, TensorStructType, TensorType
label_desc: Use image node_id 6 for calling the TorchPolicy obj's underlying member method code with example usage: obj.load_batch_into_buffer(batch, buffer_index) and returns: len, len