[bugfix] Quick fix of #1547 (#1600)

* upd

* upd

python/dgl/backend/backend.py

@@ -269,7 +269,7 @@ def asnumpy(input):
     """
     pass
 
-def copy_to(input, ctx):
+def copy_to(input, ctx, **kwargs):
     """Copy the given tensor to the context.
 
     Parameters

python/dgl/backend/mxnet/tensor.py

@@ -115,7 +115,7 @@ def astype(input, ty):
 def asnumpy(input):
     return input.asnumpy()
 
-def copy_to(input, ctx):
+def copy_to(input, ctx, **kwargs):
     return input.as_in_context(ctx)
 
 def sum(input, dim, keepdims=False):

python/dgl/backend/pytorch/tensor.py

@@ -87,13 +87,13 @@ def asnumpy(input):
     else:
         return input.cpu().detach().numpy()
 
-def copy_to(input, ctx):
+def copy_to(input, ctx, **kwargs):
     if ctx.type == 'cpu':
         return input.cpu()
     elif ctx.type == 'cuda':
         if ctx.index is not None:
             th.cuda.set_device(ctx.index)
-        return input.cuda()
+        return input.cuda(**kwargs)
     else:
         raise RuntimeError('Invalid context', ctx)
 

python/dgl/backend/tensorflow/tensor.py

@@ -129,7 +129,7 @@ def asnumpy(input):
         return input.numpy()
 
 
-def copy_to(input, ctx):
+def copy_to(input, ctx, **kwargs):
     with tf.device(ctx):
         new_tensor = tf.identity(input)
     return new_tensor

python/dgl/graph.py

@@ -3872,7 +3872,7 @@ def __repr__(self):
                           edata=str(self.edge_attr_schemes()))
 
     # pylint: disable=invalid-name
-    def to(self, ctx):
+    def to(self, ctx, **kwargs):
         """Move both ndata and edata to the targeted mode (cpu/gpu)
         Framework agnostic
 
@@ -3898,9 +3898,9 @@ def to(self, ctx):
         >>> G = G.to(torch.device('cuda:0'))
         """
         for k in self.ndata.keys():
-            self.ndata[k] = F.copy_to(self.ndata[k], ctx)
+            self.ndata[k] = F.copy_to(self.ndata[k], ctx, **kwargs)
         for k in self.edata.keys():
-            self.edata[k] = F.copy_to(self.edata[k], ctx)
+            self.edata[k] = F.copy_to(self.edata[k], ctx, **kwargs)
         return self
     # pylint: enable=invalid-name
 

python/dgl/heterograph.py

@@ -4025,7 +4025,7 @@ def filter_edges(self, predicate, edges=ALL, etype=None):
             edges = F.tensor(edges)
             return F.boolean_mask(edges, e_mask)
 
-    def to(self, ctx):  # pylint: disable=invalid-name
+    def to(self, ctx, **kwargs):  # pylint: disable=invalid-name
         """Move both ndata and edata to the targeted mode (cpu/gpu)
         Framework agnostic