From 3e19457912a5310c3f7d566b668a23f1b9f92c3c Mon Sep 17 00:00:00 2001
From: Wenyong Huang <weyo.huang@gmail.com>
Date: Wed, 1 Aug 2018 22:57:11 +0800
Subject: avoid cudaMalloc in CopyRows

---
 src/graph/node_operators_unary.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'src/graph/node_operators_unary.h')

diff --git a/src/graph/node_operators_unary.h b/src/graph/node_operators_unary.h
index cbeded24..b51c8a75 100644
--- a/src/graph/node_operators_unary.h
+++ b/src/graph/node_operators_unary.h
@@ -657,7 +657,7 @@ struct RowsNodeOp : public UnaryNodeOp {
   NodeOps forwardOps() {
     // @TODO: solve this with a tensor!
 
-    return {NodeOp(CopyRows(val_, child(0)->val(), indices_))};
+    return {NodeOp(CopyRows(val_, child(0)->val(), indices_, graph()->allocator()))};
   }
 
   NodeOps backwardOps() {
-- 
cgit v1.2.3