class_docstr: A Structure stores a Torch Tensor and labeled states.
  Inspired from the paper:
  PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
  
  https://arxiv.org/abs/2108.05818
class_docstr_tok: ["A","Structure","stores","a","Torch","Tensor","and","labeled","states",".","Inspired","from","the","paper",":","PatrickStar",":","Parallel","Training","of","Pre-trained","Models","via","Chunk-based","Memory","Management","https",":","\/\/arxiv.org\/abs\/2108.05818"]
docstr: Update global manager when changing the state of a tensor
docstr_tok: ["Update","global","manager","when","changing","the","state","of","a","tensor"]
code: def __trans_state_update(
      self, from_state: TensorState, to_state: TensorState
  ):
      """Update global manager when changing the state of a tensor"""
      manager = StatefulTensor.GST_MGR
      size = self.payload_size
      device_type = self.device.type
  
      if from_state != TensorState.FREE:
          manager.state_mem[device_type][from_state] -= size
      else:
          # when from_state is FREE, the tensor is new to manager
          # we should add its memory
          manager.total_mem[device_type] += size
  
      if to_state != TensorState.FREE:
          manager.state_mem[device_type][to_state] += size
      else:
          # when to_state is FREE, the tensor will be deleted soon
          # we should sub its memory
          manager.total_mem[device_type] -= size
  
imports: from enum import Enum
  from typing import Optional, Union
  import torch
  from .gemini_context import GeminiMemoryManager
label_desc: Use image node_id 19 for calling the StatefulTensor obj's underlying member method code with example usage: obj.__trans_state_update(from_state, to_state) without return types