File tree Expand file tree Collapse file tree 2 files changed +17
-2
lines changed
Expand file tree Collapse file tree 2 files changed +17
-2
lines changed Original file line number Diff line number Diff line change @@ -61,6 +61,22 @@ struct C10_API Storage {
6161 allocator,
6262 resizable)) {}
6363
64+ // Creates storage with pre-allocated memory buffer. Allocator is given for
65+ // potential future reallocations, however it can be nullptr if the storage
66+ // is non-resizable
67+ Storage (
68+ use_byte_size_t /* use_byte_size*/ ,
69+ SymInt size_bytes,
70+ at::DataPtr data_ptr,
71+ at::Allocator* allocator = nullptr ,
72+ bool resizable = false )
73+ : storage_impl_(c10::make_intrusive<StorageImpl>(
74+ StorageImpl::use_byte_size_t (),
75+ std::move(size_bytes),
76+ std::move(data_ptr),
77+ allocator,
78+ resizable)) {}
79+
6480 protected:
6581 explicit Storage (unsafe_borrow_t , const Storage& rhs)
6682 : storage_impl_(c10::intrusive_ptr<c10::StorageImpl>::reclaim(
Original file line number Diff line number Diff line change @@ -660,10 +660,9 @@ static Tensor make_tensor_for_subclass_helper(
660660 Storage storage{
661661 Storage::use_byte_size_t {},
662662 size_bytes,
663+ at::DataPtr{nullptr , options.device ()},
663664 /* allocator=*/ c10::GetAllocator (c10::kMeta ),
664665 /* resizable=*/ true };
665- // TODO: constructor should probably accept data pointer
666- storage.set_data_ptr_noswap (at::DataPtr{nullptr , options.device ()});
667666
668667 auto keys = c10::DispatchKeySet ({options.computeDispatchKey ()});
669668 if (extra_dispatch_keys.has_value ()) {
You can’t perform that action at this time.
0 commit comments