class_docstr: Mixin class to add support for Hugging Face accelerate
  
  This is an *experimental* feature.
  
  Use this mixin class with one of the neural net classes (e.g. ``NeuralNet``,
  ``NeuralNetClassifier``, or ``NeuralNetRegressor``) and pass an instance of
  ``Accelerator`` for mixed precision, multi-GPU, or TPU training.
  
  Install the accelerate library using:
  
  .. code-block::
  
    python -m pip install accelerate
  
  skorch does not itself provide any facilities to enable these training
  features. A lot of them can still be implemented by the user with a little
  bit of extra work but it can be a daunting task. That is why this helper
  class was added: Using this mixin in conjunction with the accelerate library
  should cover a lot of common use cases.
  
  .. note::
  
      Under the hood, accelerate uses :class:`~torch.cuda.amp.GradScaler`,
      which does not support passing the training step as a closure.
      Therefore, if your optimizer requires that (e.g.
      :class:`torch.optim.LBFGS`), you cannot use accelerate.
  
  .. warning::
  
      Since accelerate is still quite young and backwards compatiblity
      breaking features might be added, we treat its integration as an
      experimental feature. When accelerate's API stabilizes, we will consider
      adding it to skorch proper.
  
      Also, models accelerated this way cannot be pickled. If you need to save
      and load the net, either use :py:meth:`skorch.net.NeuralNet.save_params`
      and :py:meth:`skorch.net.NeuralNet.load_params` or don't use
      ``accelerate``.
  
  Examples
  --------
  >>> from skorch import NeuralNetClassifier
  >>> from skorch.hf import AccelerateMixin
  >>> from accelerate import Accelerator
  >>>
  >>> class AcceleratedNet(AccelerateMixin, NeuralNetClassifier):
  ...     '''NeuralNetClassifier with accelerate support'''
  >>>
  >>> accelerator = Accelerator(...)
  >>> # you may pass gradient_accumulation_steps to enable grad accumulation
  >>> net = AcceleratedNet(MyModule,  accelerator=accelerator)
  >>> net.fit(X, y)
  
  The same approach works with all the other skorch net classes.
  
  Parameters
  ----------
  accelerator : accelerate.Accelerator
    In addition to the usual parameters, pass an instance of
    ``accelerate.Accelerator`` with the desired settings.
  
  device : str, torch.device, or None (default=None)
    The compute device to be used. When using accelerate, it is recommended to
    leave device handling to accelerate. Therefore, it is best to leave this
    argument to be None, which means that skorch does not set the device.
  
  unwrap_after_train : bool (default=True)
    By default, with this option being ``True``, the module(s) and criterion
    are automatically "unwrapped" after training. This means that their
    initial state -- from before they were prepared by the ``accelerator`` --
    is restored. This is necessary to pickle the net.
  
    There are circumstances where you might want to disable this behavior. For
    instance, when you want to further train the model with AMP enabled (using
    ``net.partial_fit`` or ``warm_start=True``). Also, unwrapping the modules
    means that the advantage of using mixed precision is lost during
    inference. In those cases, if you don't need to pickle the net, you should
    set ``unwrap_after_train=False``.
  
  callbacks__print_log__sink : 'auto' or callable
    If 'auto', uses the ``print`` function of the accelerator, if it has one.
    This avoids printing the same output multiple times when training
    concurrently on multiple machines. If the accelerator does not have a
    ``print`` function, use Python's ``print`` function instead.
class_docstr_tok: ["Mixin","class","to","add","support","for","Hugging","Face","accelerate","This","is","an","*","experimental","*","feature",".","Use","this","mixin","class","with","one","of","the","neural","net","classes","(","e.g",".","``","NeuralNet","``",",","``","NeuralNetClassifier","``",",","or","``","NeuralNetRegressor","``",")","and","pass","an","instance","of","``","Accelerator","``","for","mixed","precision",",","multi-GPU",",","or","TPU","training",".","Install","the","accelerate","library","using",":","..","code-block",":",":","python","-m","pip","install","accelerate","skorch","does","not","itself","provide","any","facilities","to","enable","these","training","features",".","A","lot","of","them","can","still","be","implemented","by","the","user","with","a","little","bit","of","extra","work","but","it","can","be","a","daunting","task",".","That","is","why","this","helper","class","was","added",":","Using","this","mixin","in","conjunction","with","the","accelerate","library","should","cover","a","lot","of","common","use","cases",".","..","note",":",":","Under","the","hood",",","accelerate","uses",":","class",":","`","~torch.cuda.amp.GradScaler","`",",","which","does","not","support","passing","the","training","step","as","a","closure",".","Therefore",",","if","your","optimizer","requires","that","(","e.g",".",":","class",":","`","torch.optim.LBFGS","`",")",",","you","can","not","use","accelerate",".","..","warning",":",":","Since","accelerate","is","still","quite","young","and","backwards","compatiblity","breaking","features","might","be","added",",","we","treat","its","integration","as","an","experimental","feature",".","When","accelerate","'s","API","stabilizes",",","we","will","consider","adding","it","to","skorch","proper",".","Also",",","models","accelerated","this","way","can","not","be","pickled",".","If","you","need","to","save","and","load","the","net",",","either","use",":","py",":","meth",":","`","skorch.net.NeuralNet.save_params","`","and",":","py",":","meth",":","`","skorch.net.NeuralNet.load_params","`","or","do","n't","use","``","accelerate","``",".","Examples","--","--","--","--",">",">",">","from","skorch","import","NeuralNetClassifier",">",">",">","from","skorch.hf","import","AccelerateMixin",">",">",">","from","accelerate","import","Accelerator",">",">",">",">",">",">","class","AcceleratedNet","(","AccelerateMixin",",","NeuralNetClassifier",")",":","...","``","'NeuralNetClassifier","with","accelerate","support","''","'",">",">",">",">",">",">","accelerator","=","Accelerator","(","...",")",">",">",">","#","you","may","pass","gradient_accumulation_steps","to","enable","grad","accumulation",">",">",">","net","=","AcceleratedNet","(","MyModule",",","accelerator=accelerator",")",">",">",">","net.fit","(","X",",","y",")","The","same","approach","works","with","all","the","other","skorch","net","classes",".","Parameters","--","--","--","--","--","accelerator",":","accelerate.Accelerator","In","addition","to","the","usual","parameters",",","pass","an","instance","of","``","accelerate.Accelerator","``","with","the","desired","settings",".","device",":","str",",","torch.device",",","or","None","(","default=None",")","The","compute","device","to","be","used",".","When","using","accelerate",",","it","is","recommended","to","leave","device","handling","to","accelerate",".","Therefore",",","it","is","best","to","leave","this","argument","to","be","None",",","which","means","that","skorch","does","not","set","the","device",".","unwrap_after_train",":","bool","(","default=True",")","By","default",",","with","this","option","being","``","True","``",",","the","module","(","s",")","and","criterion","are","automatically","``","unwrapped","''","after","training",".","This","means","that","their","initial","state","--","from","before","they","were","prepared","by","the","``","accelerator","``","--","is","restored",".","This","is","necessary","to","pickle","the","net",".","There","are","circumstances","where","you","might","want","to","disable","this","behavior",".","For","instance",",","when","you","want","to","further","train","the","model","with","AMP","enabled","(","using","``","net.partial_fit","``","or","``","warm_start=True","``",")",".","Also",",","unwrapping","the","modules","means","that","the","advantage","of","using","mixed","precision","is","lost","during","inference",".","In","those","cases",",","if","you","do","n't","need","to","pickle","the","net",",","you","should","set","``","unwrap_after_train=False","``",".","callbacks__print_log__sink",":","'auto","'","or","callable","If","'auto","'",",","uses","the","``","print","``","function","of","the","accelerator",",","if","it","has","one",".","This","avoids","printing","the","same","output","multiple","times","when","training","concurrently","on","multiple","machines",".","If","the","accelerator","does","not","have","a","``","print","``","function",",","use","Python","'s","``","print","``","function","instead","."]
code: def save_params(self, *args, **kwargs):
      # has to be called even if not main process, or else there is a dead lock
      self.accelerator.wait_for_everyone()
  
      if not self._wrapped_with_accelerator:
          if self.accelerator.is_main_process:
              super().save_params(*args, **kwargs)
      else:
          # A potential issue with using accelerate is that a model that has
          # been prepared with accelerate is wrapped, so that the keys of the
          # state dict have an additional prefix, "module.". Therefore, when
          # the model is unwrapped when saving and wrapped when loading, or
          # vice versa, there will be a mismatch in the state dict keys. To
          # prevent this, always unwrap before saving. During loading, in case
          # the model is wrapped, this would result in an error, but we take
          # care of unwrapping the model in that case during loading.
          self._unwrap_accelerator()
          try:
              # note: although saving is only done on the main process,
              # unwrapping+wrapping has to be done on all processes, or else
              # there is an error, not sure why
              if self.accelerator.is_main_process:
                  super().save_params(*args, **kwargs)
          finally:
              self._initialize_accelerator()
  
imports: import io
  import os
  import pathlib
  from copy import deepcopy
  from operator import itemgetter
  import numpy
  import torch
  from sklearn.base import BaseEstimator, TransformerMixin
  from skorch.callbacks import LRScheduler
  from skorch.dataset import unpack_data
  from skorch.utils import check_is_fitted, params_for
label_desc: Use image node_id 13 for calling the AccelerateMixin obj's underlying member method code with example usage: obj.save_params() without return types