class_docstr: Interface between a PyTorch model and a Thinc Model. This container is
  *not* a Thinc Model subclass itself.
  
  mixed_precision:
      Enable mixed-precision. This changes whitelisted ops to run
      in half precision for better performance and lower memory use.
  grad_scaler:
      The gradient scaler to use for mixed-precision training. If this
      argument is set to "None" and mixed precision is enabled, a gradient
      scaler with the default configuration is used.
  device:
      The PyTorch device to run the model on. When this argument is
      set to "None", the default device for the currently active Thinc
      ops is used.
  serialize_model:
      Callback that receives the wrapped PyTorch model as its argument and
      returns a "bytes" representation of the same. The representation should
      contain all the necessary information to fully deserialize the model.
  deserialize_model:
      Callback that receives the default PyTorch model (passed to the constructor), the
      serialized "bytes" representation and a PyTorch device. It should return a
      fully deserialized model on the target device as its result.
class_docstr_tok: ["Interface","between","a","PyTorch","model","and","a","Thinc","Model",".","This","container","is","*","not","*","a","Thinc","Model","subclass","itself",".","mixed_precision",":","Enable","mixed-precision",".","This","changes","whitelisted","ops","to","run","in","half","precision","for","better","performance","and","lower","memory","use",".","grad_scaler",":","The","gradient","scaler","to","use","for","mixed-precision","training",".","If","this","argument","is","set","to","``","None","''","and","mixed","precision","is","enabled",",","a","gradient","scaler","with","the","default","configuration","is","used",".","device",":","The","PyTorch","device","to","run","the","model","on",".","When","this","argument","is","set","to","``","None","''",",","the","default","device","for","the","currently","active","Thinc","ops","is","used",".","serialize_model",":","Callback","that","receives","the","wrapped","PyTorch","model","as","its","argument","and","returns","a","``","bytes","''","representation","of","the","same",".","The","representation","should","contain","all","the","necessary","information","to","fully","deserialize","the","model",".","deserialize_model",":","Callback","that","receives","the","default","PyTorch","model","(","passed","to","the","constructor",")",",","the","serialized","``","bytes","''","representation","and","a","PyTorch","device",".","It","should","return","a","fully","deserialized","model","on","the","target","device","as","its","result","."]
docstr: Pass the inputs through to the underlying PyTorch model, keeping
  track of which items in the input are tensors requiring gradients.
  If the model returns a single value, it is converted into a one-element tuple.
  Return the outputs and a callback to backpropagate.
docstr_tok: ["Pass","the","inputs","through","to","the","underlying","PyTorch","model",",","keeping","track","of","which","items","in","the","input","are","tensors","requiring","gradients",".","If","the","model","returns","a","single","value",",","it","is","converted","into","a","one-element","tuple",".","Return","the","outputs","and","a","callback","to","backpropagate","."]
code: def begin_update(self, inputs: ArgsKwargs):
      """Pass the inputs through to the underlying PyTorch model, keeping
      track of which items in the input are tensors requiring gradients.
      If the model returns a single value, it is converted into a one-element tuple.
      Return the outputs and a callback to backpropagate.
      """
      self._model.train()
  
      # Note: mixed-precision autocast must not be applied to backprop.
      with torch.cuda.amp.autocast(self._mixed_precision):
          output = self._model(*inputs.args, **inputs.kwargs)
  
      def backprop(grads):
          # Normally, gradient scaling is applied to the loss of a model. However,
          # since regular thinc layers do not use mixed-precision, we perform scaling
          # locally in this shim. Scaling the loss by a factor, scales the gradients
          # by the same factor (see the chain rule). Therefore, we scale the gradients
          # backprop'ed through the succeeding layer to get the same effect as loss
          # scaling.
          grads.kwargs["grad_tensors"] = self._grad_scaler.scale(
              grads.kwargs["grad_tensors"], inplace=True
          )
  
          torch.autograd.backward(*grads.args, **grads.kwargs)
  
          # Unscale weights and check for overflows during backprop.
          grad_tensors = []
          for torch_data in itertools.chain(
              self._model.parameters(),
              iterate_recursive(lambda x: hasattr(x, "grad"), inputs),
          ):
              if torch_data.grad is not None:
                  grad_tensors.append(torch_data.grad)
          found_inf = self._grad_scaler.unscale(grad_tensors)
  
          # If there was an over/underflow, return zeroed-out gradients.
          if found_inf:
              grad_get = (
                  lambda x: x.grad.zero_()
                  if x.grad is not None
                  else x.grad
              )
          else:
              grad_get = lambda x: x.grad
  
          return convert_recursive(
              lambda x: hasattr(x, "grad"), grad_get, inputs
          )
  
      return output, backprop
  
imports: import contextlib
  import itertools
  from io import BytesIO
  from typing import Any, Callable, Dict, Optional, cast
  import srsly
  from ..backends import CupyOps, context_pools, get_current_ops, set_gpu_allocator
  from ..compat import torch
  from ..optimizers import Optimizer
  from ..types import ArgsKwargs, FloatsXd
  from ..util import convert_recursive, get_torch_default_device, iterate_recursive, torch2xp, xp2torch
  from .pytorch_grad_scaler import PyTorchGradScaler
  from .shim import Shim
label_desc: Use image node_id 5 for calling the PyTorchShim obj's underlying member method code with example usage: obj.begin_update(inputs) and returns: output,  backprop, convert_recursive