[caffe2] Bypass memonger for in-place ops (#46378)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/46378
Reviewed By: dzhulgakov
Differential Revision: D24236604
fbshipit-source-id: 9f599687467ea969e89243482f8e2a41f7db0a23
diff --git a/caffe2/core/memonger.cc b/caffe2/core/memonger.cc
index 994c97c..8833a0c 100644
--- a/caffe2/core/memonger.cc
+++ b/caffe2/core/memonger.cc
@@ -24,6 +24,25 @@
LOG(INFO) << "Memonger does not support RecurrentNetwork yet";
return net;
}
+
+ // Check for in-place ops
+ auto* schema = OpSchemaRegistry::Schema(op.type());
+ if (schema) {
+ // Memonger modifies the graph. Do an early schema check here to make sure
+ // the operators are valid
+ CAFFE_ENFORCE(
+ schema->Verify(op),
+ "Operator def did not pass schema checking: ",
+ ProtoDebugString(op));
+ for (int in_idx = 0; in_idx < op.input_size(); in_idx++) {
+ for (int out_idx = 0; out_idx < op.output_size(); out_idx++) {
+ if (schema->inplace_enforced(in_idx, out_idx)) {
+ LOG(INFO) << "Memonger does not support in-place ops yet";
+ return net;
+ }
+ }
+ }
+ }
ops.push_back(op);
}
