frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use CPU Allocator for reading from zip container

Browse Source 
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
This commit is contained in:
David Reiss
2020-06-25 10:44:04 -07:00
committed by Facebook GitHub Bot
parent c362138f43
commit 7369dc8d1f
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
caffe2/serialize/inline_container.cc
View File

@ -7,6 +7,7 @@
#include <algorithm>
#include <c10/core/Allocator.h>
#include <c10/core/CPUAllocator.h>
#include <c10/core/Backend.h>
#include "caffe2/core/common.h"
@ -218,11 +219,10 @@ std::tuple<at::DataPtr, size_t> 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);
}