Method alloc

csrc/aio/py_lib/deepspeed_pin_tensor.cpp:23–34

}

torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem,
                                            const torch::TensorOptions& options)
{
    const auto scalar_dtype = torch::typeMetaToScalarType(options.dtype());
    const auto num_bytes = num_elem * torch::elementSize(scalar_dtype);
    auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true);
    assert(nullptr != pinned_buffer);

    _locked_tensors[pinned_buffer] = num_bytes;

    return at::from_blob(pinned_buffer, static_cast<int64_t>(num_elem), options);
}

torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type)
{