forked from facebookresearch/torchdim
-
Notifications
You must be signed in to change notification settings - Fork 0
/
functorch.diff
34 lines (32 loc) · 1.13 KB
/
functorch.diff
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
34
diff --git a/functorch/csrc/BatchedTensorImpl.h b/functorch/csrc/BatchedTensorImpl.h
index 45999e0..3610dc3 100644
--- a/functorch/csrc/BatchedTensorImpl.h
+++ b/functorch/csrc/BatchedTensorImpl.h
@@ -73,6 +73,9 @@ struct BatchedTensorImpl : public c10::TensorImpl {
void set_size(int64_t dim, int64_t new_size) override;
void set_stride(int64_t dim, int64_t new_stride) override;
void set_storage_offset(int64_t storage_offset) override;
+ void set_level(int64_t level) {
+ level_ = level;
+ }
#ifdef DEBUG
bool has_storage() const override;
#endif
diff --git a/functorch/csrc/init.cpp b/functorch/csrc/init.cpp
index c09a53d..371b6fc 100644
--- a/functorch/csrc/init.cpp
+++ b/functorch/csrc/init.cpp
@@ -358,7 +358,7 @@ static void dump_local_tls() {
namespace at { namespace functorch {
-
+#if 0
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("_add_batch_dim", &at::functorch::_add_batch_dim, "add batch dim");
m.def("_remove_batch_dim", &at::functorch::_remove_batch_dim, "remove batch dim");
@@ -410,5 +410,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
initDispatchBindings(m.ptr());
#endif
}
-
+#endif
}}