Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "213e27f05c9f1249837140affca955f6ac37de80"
Unverified Commit 6de7d5fa authored by Hao Yuan's avatar Hao Yuan Committed by GitHub
Browse files

[Bugfix] Fix conditional judgment of cuda (#4026) (#4027)

parent 4b87e47f
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
examples/pytorch/gat/train.py
View file @ 6de7d5fa
...@@ -107,6 +107,7 @@ def main(args): ...@@ -107,6 +107,7 @@ def main(args):
for epoch in range(args.epochs): for epoch in range(args.epochs):
model.train() model.train()
if epoch >= 3: if epoch >= 3:
if cuda:
torch.cuda.synchronize() torch.cuda.synchronize()
t0 = time.time() t0 = time.time()
# forward # forward
...@@ -118,6 +119,7 @@ def main(args): ...@@ -118,6 +119,7 @@ def main(args):
optimizer.step() optimizer.step()
if epoch >= 3: if epoch >= 3:
if cuda:
torch.cuda.synchronize() torch.cuda.synchronize()
dur.append(time.time() - t0) dur.append(time.time() - t0)
......
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