class_docstr: A training loop.
class_docstr_tok: ["A","training","loop","."]
docstr: Find the allowable sub batch size for training with the current setting.
  
  This method checks if it is possible to train the model with the given training data and the desired batch size
  on the hardware at hand. If possible, the sub-batch size equals the batch size. Otherwise, the maximum
  permissible sub-batch size is determined.
  
  :param batch_size:
      The initial batch size to start with.
  :param sampler:
      The sampler (None or schlichtkrull)
  :param triples_factory:
      A triples factory
  
  :return:
      Tuple containing the sub-batch size to use and indicating if the search was finished, i.e. successfully
      without hardware errors, as well as if sub-batching is possible
  
  :raises RuntimeError:
      If a runtime error is raised during training
docstr_tok: ["Find","the","allowable","sub","batch","size","for","training","with","the","current","setting",".","This","method","checks","if","it","is","possible","to","train","the","model","with","the","given","training","data","and","the","desired","batch","size","on","the","hardware","at","hand",".","If","possible",",","the","sub-batch","size","equals","the","batch","size",".","Otherwise",",","the","maximum","permissible","sub-batch","size","is","determined",".",":","param","batch_size",":","The","initial","batch","size","to","start","with",".",":","param","sampler",":","The","sampler","(","None","or","schlichtkrull",")",":","param","triples_factory",":","A","triples","factory",":","return",":","Tuple","containing","the","sub-batch","size","to","use","and","indicating","if","the","search","was","finished",",","i.e",".","successfully","without","hardware","errors",",","as","well","as","if","sub-batching","is","possible",":","raises","RuntimeError",":","If","a","runtime","error","is","raised","during","training"]
code: def _sub_batch_size_search(
      self,
      *,
      batch_size: int,
      sampler: Optional[str],
      triples_factory: CoreTriplesFactory,
  ) -> Tuple[int, bool, bool]:
      """Find the allowable sub batch size for training with the current setting.
  
      This method checks if it is possible to train the model with the given training data and the desired batch size
      on the hardware at hand. If possible, the sub-batch size equals the batch size. Otherwise, the maximum
      permissible sub-batch size is determined.
  
      :param batch_size:
          The initial batch size to start with.
      :param sampler:
          The sampler (None or schlichtkrull)
      :param triples_factory:
          A triples factory
  
      :return:
          Tuple containing the sub-batch size to use and indicating if the search was finished, i.e. successfully
          without hardware errors, as well as if sub-batching is possible
  
      :raises RuntimeError:
          If a runtime error is raised during training
      """
      sub_batch_size = batch_size
      finished_search = False
      supports_sub_batching = True
  
      try:
          # The cache of the previous run has to be freed to allow accurate memory availability estimates
          self._free_graph_and_cache()
          logger.debug(
              f"Trying batch_size {batch_size} for training now."
          )
          self._train(
              triples_factory=triples_factory,
              num_epochs=1,
              batch_size=batch_size,
              sub_batch_size=sub_batch_size,
              sampler=sampler,
              only_size_probing=True,
          )
      except RuntimeError as runtime_error:
          self._free_graph_and_cache()
          if not is_cudnn_error(
              runtime_error
          ) and not is_cuda_oom_error(runtime_error):
              raise runtime_error
          logger.debug(
              f"The batch_size {batch_size} was too big, sub_batching is required."
          )
          sub_batch_size //= 2
      else:
          finished_search = True
          logger.debug("No sub-batching required.")
  
      if not finished_search:
          logger.info(
              "Starting sub_batch_size search for training now..."
          )
          if get_batchnorm_modules(
              self.model
          ):  # if there are any, this is truthy
              logger.info("This model does not support sub-batching.")
              supports_sub_batching = False
              sub_batch_size = batch_size
          else:
              while True:
                  logger.debug(
                      f"Trying sub_batch_size {sub_batch_size} now."
                  )
                  try:
                      self._free_graph_and_cache()
                      self._train(
                          num_epochs=1,
                          batch_size=batch_size,
                          sub_batch_size=sub_batch_size,
                          sampler=sampler,
                          only_size_probing=True,
                          triples_factory=triples_factory,
                      )
                  except RuntimeError as runtime_error:
                      self._free_graph_and_cache()
                      if not is_cudnn_error(
                          runtime_error
                      ) and not is_cuda_oom_error(runtime_error):
                          raise runtime_error
                      if sub_batch_size == 1:
                          logger.info(
                              f"Even sub_batch_size={sub_batch_size} does not fit in memory with these parameters",
                          )
                          break
                      logger.debug(
                          f"The sub_batch_size {sub_batch_size} was too big, trying less now."
                      )
                      sub_batch_size //= 2
                  else:
                      finished_search = True
                      logger.info(
                          f"Concluded search with sub_batch_size {sub_batch_size}."
                      )
                      break
  
      self._free_graph_and_cache()
  
      return sub_batch_size, finished_search, supports_sub_batching
  
imports: import gc
  import logging
  import os
  import pathlib
  import pickle
  import random
  import time
  from abc import ABC, abstractmethod
  from contextlib import ExitStack
  from datetime import datetime
  from hashlib import md5
  from tempfile import NamedTemporaryFile, TemporaryDirectory
  from typing import IO, Any, ClassVar, Generic, List, Mapping, Optional, Tuple, TypeVar, Union
  import numpy
  import torch
  from class_resolver import HintOrType, OptionalKwargs
  from class_resolver.contrib.torch import lr_scheduler_resolver, optimizer_resolver
  from torch.optim.optimizer import Optimizer
  from torch.utils.data import DataLoader
  from tqdm.autonotebook import tqdm, trange
  from .callbacks import GradientAbsClippingTrainingCallback, GradientNormClippingTrainingCallback, LearningRateSchedulerTrainingCallback, MultiTrainingCallback, OptimizerTrainingCallback, StopperTrainingCallback, TrackerTrainingCallback, TrainingCallback, TrainingCallbackHint, TrainingCallbackKwargsHint
  from ..constants import PYKEEN_CHECKPOINTS, PYKEEN_DEFAULT_CHECKPOINT
  from ..lr_schedulers import LRScheduler
  from ..models import RGCN, Model
  from ..stoppers import Stopper
  from ..trackers import ResultTracker, tracker_resolver
  from ..triples import CoreTriplesFactory, TriplesFactory
  from ..typing import InductiveMode
  from ..utils import format_relative_comparison, get_batchnorm_modules, get_preferred_device, is_cuda_oom_error, is_cudnn_error, normalize_string
label_desc: Use image node_id 16 for calling the TrainingLoop obj's underlying member method code with example usage: obj._sub_batch_size_search() and returns: sub_batch_size,  finished_search,  supports_sub_batching