mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
b7ebc00979
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/11924 Previous diffs removed Blob -> caffe2 dependencies, now we can move it to ATen/core. This is pre-work for allowing storing Blob in IValue. Reviewed By: ezyang Differential Revision: D9980641 fbshipit-source-id: 32082a673ec94c42c20b2298adced8bb7ca94d07
44 lines
1.1 KiB
C++
44 lines
1.1 KiB
C++
#ifndef CAFFE2_CORE_BLOB_H_
|
|
#define CAFFE2_CORE_BLOB_H_
|
|
|
|
#include <cstddef>
|
|
#include <sstream>
|
|
#include <typeinfo>
|
|
#include <type_traits>
|
|
#include <vector>
|
|
#include "caffe2/core/common.h"
|
|
|
|
#include <ATen/core/blob.h>
|
|
#include <ATen/core/typeid.h>
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/tensor.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
inline bool BlobIsTensorType(const Blob& blob, DeviceType device_type) {
|
|
bool is_match = blob.meta().Match<Tensor>();
|
|
if (!is_match) {
|
|
return false;
|
|
}
|
|
const Tensor* tensor = &blob.Get<Tensor>();
|
|
return tensor && tensor->GetDeviceType() == device_type;
|
|
}
|
|
|
|
inline Tensor* BlobGetMutableTensor(Blob* blob, DeviceType device_type) {
|
|
if (blob->IsType<Tensor>()) {
|
|
Tensor* tensor = blob->GetMutable<Tensor>();
|
|
if (tensor->GetDeviceType() == device_type) {
|
|
return tensor;
|
|
}
|
|
}
|
|
|
|
// if we're here, then either Blob didn't hold a Tensor
|
|
// or that Tensor had the wrong DeviceType.
|
|
VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<Tensor>()
|
|
<< " DeviceType:" << device_type;
|
|
return blob->Reset<Tensor>(new Tensor(device_type));
|
|
}
|
|
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_CORE_BLOB_H_
|