diff options
| author | Trevor Killeen <killeentm@gmail.com> | 2017-07-10 19:42:21 +0300 |
|---|---|---|
| committer | Soumith Chintala <soumith@gmail.com> | 2017-07-13 07:37:13 +0300 |
| commit | 2186e414ad8fc4dfc9f2ed090c0bf8a0e1946e62 (patch) | |
| tree | a42868e8ce517986f783e16fbd8f29e2832cab5b | |
| parent | eda1c8892fdb054bba0869bc81355eaad5260c2c (diff) | |
Advanced Indexing: Calculate linear offsets directly on the GPU when working with CUDA Tensors
| -rw-r--r-- | lib/TH/generic/THTensor.c | 8 |
1 files changed, 6 insertions, 2 deletions
diff --git a/lib/TH/generic/THTensor.c b/lib/TH/generic/THTensor.c index e44e06e..8ebec67 100644 --- a/lib/TH/generic/THTensor.c +++ b/lib/TH/generic/THTensor.c @@ -321,8 +321,8 @@ void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count) { THArgCheck(THTensor_(nDimension)(ops[i]) > 0, i, "can't expand empty tensor %d", i); } - long *op_sizes[count]; - long op_dims[count]; + long **op_sizes = THAlloc(sizeof(long*) * count); + long *op_dims = THAlloc(sizeof(long) * count); for (int i = 0; i < count; ++i) { op_sizes[i] = ops[i]->size; @@ -339,6 +339,8 @@ void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count) { 1024); if(ret != 0) { + THFree(op_sizes); + THFree(op_dims); THLongStorage_free(sizes); THError(error_buffer); return; @@ -348,6 +350,8 @@ void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count) { THTensor_(expand)(rets[i], ops[i], sizes); } + THFree(op_sizes); + THFree(op_dims); THLongStorage_free(sizes); } |