Skip to content

Commit a63221a

Browse files
swolchokpytorchmergebot
authored andcommitted
Fix TODO in make_tensor_for_subclass_helper (pytorch#162336)
The constructor does accept a DataPtr (had to fix the DataPtr variant not accepting a SymInt, though). Pull Request resolved: pytorch#162336 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#162298
1 parent c9485f8 commit a63221a

File tree

2 files changed

+17
-2
lines changed

2 files changed

+17
-2
lines changed

c10/core/Storage.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff 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(

torch/csrc/autograd/python_variable.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff 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()) {

0 commit comments

Comments
 (0)