Skip to content

GitLab

Unverified Commit ee91863e authored by Wey Gu's avatar Wey Gu Committed by GitHub
Browse files

[Example] Fixed device type in GraphSAGE inference (#4306) 

* Fixed device type in inference

* change buffer_device instead
parent 2cf05c53
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
examples/pytorch/graphsage/node_classification.py
View file @ ee91863e
...@@ -39,7 +39,7 @@ class SAGE(nn.Module): ...@@ -39,7 +39,7 @@ class SAGE(nn.Module):
g, torch.arange(g.num_nodes()).to(g.device), sampler, device=device, g, torch.arange(g.num_nodes()).to(g.device), sampler, device=device,
batch_size=batch_size, shuffle=False, drop_last=False, batch_size=batch_size, shuffle=False, drop_last=False,
num_workers=0) num_workers=0)
buffer_device = 'cpu' buffer_device = torch.device('cpu')
pin_memory = (buffer_device != device) pin_memory = (buffer_device != device)
for l, layer in enumerate(self.layers): for l, layer in enumerate(self.layers):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment