summaryrefslogtreecommitdiff
path: root/aten
diff options
context:
space:
mode:
authorSoumith Chintala <soumith@gmail.com>2018-04-23 22:20:03 -0400
committerGitHub <noreply@github.com>2018-04-23 22:20:03 -0400
commitb6ed729cdcb7fec75978c42538f99f46f5732177 (patch)
tree21ef85bbcab72e2ae63809020a33d87cd3047b49 /aten
parentdf2817d3b129d868bfb57d40ec63a6612dbd564a (diff)
downloadpytorch-b6ed729cdcb7fec75978c42538f99f46f5732177.tar.gz
pytorch-b6ed729cdcb7fec75978c42538f99f46f5732177.tar.bz2
pytorch-b6ed729cdcb7fec75978c42538f99f46f5732177.zip
fix memory leak in median (#6889)
Diffstat (limited to 'aten')
-rw-r--r--aten/src/THC/generic/THCTensorMathReduce.cu3
1 files changed, 3 insertions, 0 deletions
diff --git a/aten/src/THC/generic/THCTensorMathReduce.cu b/aten/src/THC/generic/THCTensorMathReduce.cu
index 69b57d6ded..8dd59366f1 100644
--- a/aten/src/THC/generic/THCTensorMathReduce.cu
+++ b/aten/src/THC/generic/THCTensorMathReduce.cu
@@ -392,6 +392,9 @@ THCTensor_(median)(THCState *state,
THCTensor *newValues = THCTensor_(newNarrow)(state, sorted, dimension, k, 1);
THCudaLongTensor *newIndices = THCudaLongTensor_newNarrow(state, sorted_indices, dimension, k, 1);
+ THCTensor_(free)(state, sorted);
+ THCudaLongTensor_free(state, sorted_indices);
+
if (!keepdim) {
THCTensor_(squeeze1d)(state, newValues, newValues, dimension);
THCudaLongTensor_squeeze1d(state, newIndices, newIndices, dimension);