code: def setup_grads_device(
      self,
      params: List[torch.Tensor],
      grads_device_map: Dict[torch.Tensor, torch.device],
  ) -> None:
      raise NotImplementedError
  
imports: import functools
  import warnings
  from abc import ABC, abstractmethod
  from time import time
  from typing import Dict, List, Optional, Tuple, Type
  import torch
  from colossalai.utils import get_current_device
  from colossalai.utils.memory import colo_device_memory_capacity
  from colossalai.zero.gemini.chunk import Chunk
  from .chunk import Chunk, ChunkManager
  from .memory_tracer import ChunkMemStatsCollector
label_desc: Use image node_id 3 for calling the PlacementPolicy obj's underlying member method code with example usage: obj.setup_grads_device(params, grads_device_map) without return types