forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
IDeepRegistration.cpp
33 lines (25 loc) · 797 Bytes
/
IDeepRegistration.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/Config.h>
#include <c10/core/Allocator.h>
#if AT_MKLDNN_ENABLED()
// needs to be included only once in library.
#include <ideep_pin_singletons.hpp>
using namespace ideep;
RegisterEngineAllocator cpu_alloc(
engine::cpu_engine(),
[](size_t size) {
return c10::GetAllocator(c10::DeviceType::CPU)->raw_allocate(size);
},
[](void* p) {
c10::GetAllocator(c10::DeviceType::CPU)->raw_deallocate(p);
}
);
namespace at::native::mkldnn{
void clear_computation_cache();
void clear_computation_cache() {
// Reset computation_cache for forward convolutions
// As it also caches max number of OpenMP workers
ideep::convolution_forward::t_store().clear();
}
} // namespace at::native::mkldnn
#endif // AT_MKLDNN_ENABLED()