class_docstr: A training loop.
class_docstr_tok: ["A","training","loop","."]
docstr: Find the maximum batch size for training with the current setting.
  
  This method checks how big the batch size can be for the current model with the given training data and the
  hardware at hand. If possible, the method will output the determined batch size and a boolean value indicating
  that this batch size was successfully evaluated. Otherwise, the output will be batch size 1 and the boolean
  value will be False.
  
  :param triples_factory:
      The triples factory over which search is run
  :param batch_size:
      The batch size to start the search with. If None, set batch_size=num_triples (i.e. full batch training).
  
  :return:
      Tuple containing the maximum possible batch size as well as an indicator if the evaluation with that size
      was successful.
  
  :raises RuntimeError:
      If a runtime error is raised during training
docstr_tok: ["Find","the","maximum","batch","size","for","training","with","the","current","setting",".","This","method","checks","how","big","the","batch","size","can","be","for","the","current","model","with","the","given","training","data","and","the","hardware","at","hand",".","If","possible",",","the","method","will","output","the","determined","batch","size","and","a","boolean","value","indicating","that","this","batch","size","was","successfully","evaluated",".","Otherwise",",","the","output","will","be","batch","size","1","and","the","boolean","value","will","be","False",".",":","param","triples_factory",":","The","triples","factory","over","which","search","is","run",":","param","batch_size",":","The","batch","size","to","start","the","search","with",".","If","None",",","set","batch_size=num_triples","(","i.e",".","full","batch","training",")",".",":","return",":","Tuple","containing","the","maximum","possible","batch","size","as","well","as","an","indicator","if","the","evaluation","with","that","size","was","successful",".",":","raises","RuntimeError",":","If","a","runtime","error","is","raised","during","training"]
code: def batch_size_search(
      self,
      *,
      triples_factory: CoreTriplesFactory,
      batch_size: Optional[int] = None,
  ) -> Tuple[int, bool]:
      """Find the maximum batch size for training with the current setting.
  
      This method checks how big the batch size can be for the current model with the given training data and the
      hardware at hand. If possible, the method will output the determined batch size and a boolean value indicating
      that this batch size was successfully evaluated. Otherwise, the output will be batch size 1 and the boolean
      value will be False.
  
      :param triples_factory:
          The triples factory over which search is run
      :param batch_size:
          The batch size to start the search with. If None, set batch_size=num_triples (i.e. full batch training).
  
      :return:
          Tuple containing the maximum possible batch size as well as an indicator if the evaluation with that size
          was successful.
  
      :raises RuntimeError:
          If a runtime error is raised during training
      """
      if batch_size is None:
          batch_size = 8192
  
      # Set upper bound
      batch_size = min(batch_size, triples_factory.num_triples)
  
      reached_max = False
      evaluated_once = False
      logger.info("Starting batch_size search for training now...")
      while True:
          logger.debug(f"Trying batch_size={batch_size}.")
          try:
              self._free_graph_and_cache()
              self._train(
                  num_epochs=1,
                  batch_size=batch_size,
                  sub_batch_size=None,
                  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 batch_size == 1:
                  logger.debug(
                      f"batch_size={batch_size} does not fit into your memory with these parameters."
                  )
                  break
  
              reached_max = True
              batch_size //= 2
  
              if evaluated_once:
                  logger.info(
                      f"Concluded batch_size search with batch_size={batch_size}."
                  )
                  break
  
              logger.debug(
                  f"batch_size={batch_size} was too big, trying less now."
              )
          else:
              self._free_graph_and_cache()
              if (
                  not reached_max
                  and batch_size <= triples_factory.num_triples
              ):
                  batch_size *= 2
              else:
                  logger.info(
                      f"Concluded batch_size search with batch_size={batch_size}."
                  )
                  evaluated_once = True
                  break
  
      return batch_size, evaluated_once
  
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 13 for calling the TrainingLoop obj's underlying member method code with example usage: obj.batch_size_search() and returns: batch_size,  evaluated_once