text: Title: Trainer with GPU based Model fails while creating Masks
  Body: ### Bug with GPU Model
  Currently, while using pruning methods like `TaylorFOWeight` Pruner, If I use a model on GPU for getting the metrics (as calculated for getting masks), it fails on [line](https://github.com/microsoft/nni/blob/cbac2c5c0f7606aca8ccf08fbd418ffe3adfe427/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py#L32) while creating masks. The reason why it fails is,
  ```
  metrics - Dict[str, Tensor] where the tensor is on cuda
  wrapper.weight_mask - Tensor on cpu
  ```
  
  **What would you like to be added**:
  ```python
   # something like this, not necessarily the right answer
  metric = metric.to(wrapper.weight_mask.device)
  ```
  
  **Why is this needed**:
  Allow us to use GPU based fitting (therefore much faster) while calculating metrics for generating masks.
  
  **Without this feature, how does current nni work**：
  Fails when model is gpu based
  
  **Components that may involve changes**:
  All Sparsity Allocators under `nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py`