Let PS_weights with same device as parameter

bluefog/torch/optimizers.py

@@ -476,7 +476,7 @@ def _register_window(self):
                     raise KeyError(
                         "Cannot find parameter {} in the _parameter_names dictionary".format(name))
 
-                ps_weights = torch.Tensor([1.0]).to(p.data.dtype)
+                ps_weights = torch.Tensor([1.0]).to(p.data.dtype).to(p.data.device)
                 self._named_ps_weights[name] = ps_weights
                 # If do not modify in the C level, it is inevitable to copy
                 # the parameter once in the cat ops.