Skip to content

GitLab

Commit d772d390 authored by Minjie Wang's avatar Minjie Wang
Browse files

fix

parent a8a4fcba
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
CMakeLists.txt
View file @ d772d390
######################################## ########################################
# Borrowed and adapted from TVM project # Borrowed and adapted from TVM project
######################################## ########################################
cmake_minimum_required(VERSION 3.2) cmake_minimum_required(VERSION 2.8)
project(dgl C CXX) project(dgl C CXX)
# Utility functions # Utility functions
......
examples/pytorch/tree_lstm/train.py
View file @ d772d390
...@@ -37,7 +37,7 @@ def main(args): ...@@ -37,7 +37,7 @@ def main(args):
bg = dgl.batch(trees) bg = dgl.batch(trees)
if cuda: if cuda:
reprs = bg.get_n_repr() reprs = bg.get_n_repr()
reprs = {key : reprs[key].cuda()} reprs = {key : val.cuda() for key, val in reprs.items()}
bg.set_n_repr(reprs) bg.set_n_repr(reprs)
return bg return bg
trainset = data.SST() trainset = data.SST()
......
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