From 7369dc8d1ff1c6187b12d23514a643408c9e2e95 Mon Sep 17 00:00:00 2001 From: David Reiss Date: Thu, 25 Jun 2020 10:44:04 -0700 Subject: [PATCH] Use CPU Allocator for reading from zip container Summary: This code path is used to read tensor bodies, so we need it to respect alignment and padding requirements. Test Plan: Ran an internal test that was failing. Reviewed By: zdevito Differential Revision: D22225622 fbshipit-source-id: f2126727f96616366850642045ab9704f3885824 --- caffe2/serialize/inline_container.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/caffe2/serialize/inline_container.cc b/caffe2/serialize/inline_container.cc index 419105c12d52..eaa5b3501149 100644 --- a/caffe2/serialize/inline_container.cc +++ b/caffe2/serialize/inline_container.cc @@ -7,6 +7,7 @@ #include #include +#include #include #include "caffe2/core/common.h" @@ -218,11 +219,10 @@ std::tuple PyTorchStreamReader::getRecord(const std::string mz_zip_archive_file_stat stat; mz_zip_reader_file_stat(ar_.get(), key, &stat); valid("retrieving file meta-data for ", name.c_str()); - void * ptr = malloc(stat.m_uncomp_size); - mz_zip_reader_extract_to_mem(ar_.get(), key, ptr, stat.m_uncomp_size, 0); + at::DataPtr retval = c10::GetCPUAllocator()->allocate(stat.m_uncomp_size); + mz_zip_reader_extract_to_mem(ar_.get(), key, retval.get(), stat.m_uncomp_size, 0); valid("reading file ", name.c_str()); - at::DataPtr retval(ptr, ptr, free, at::kCPU); return std::make_tuple(std::move(retval), stat.m_uncomp_size); }