Fix WeakIValueEq (#21891)

James Reed · facebook-github-bot · commit 28ecc104f481 · 2019-06-18T10:35:35.000-07:00
Summary: Pull Request resolved: #21891 ghimport-source-id: a037850 Test Plan: Imported from OSS Differential Revision: D15871588 Pulled By: jamesr66a fbshipit-source-id: ecfdece1285c0737d0b1dc2afe959c43d9413001
diff --git a/aten/src/ATen/core/ivalue.h b/aten/src/ATen/core/ivalue.h
@@ -495,6 +495,11 @@ struct CAFFE2_API WeakIValue final {
     std::swap(tag, rhs.tag);
   }
 
+  bool isSameIdentity(const WeakIValue& rhs) const {
+    return payload.as_int == rhs.payload.as_int && tag == rhs.tag &&
+        is_intrusive_ptr == rhs.is_intrusive_ptr;
+  }
+
   IValue lock() const {
     if (!is_intrusive_ptr) {
       return IValue(payload, tag, false);
diff --git a/torch/csrc/jit/tracer.h b/torch/csrc/jit/tracer.h
@@ -75,7 +75,7 @@ struct TORCH_API TracingState
 
   struct WeakIValueEq {
     bool operator()(const WeakIValue& t1, const WeakIValue& t2) const {
-      return t1.lock().isSameIdentity(t2.lock());
+      return t1.isSameIdentity(t2);
     }
   };
 

Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ struct TORCH_API TracingState`
`75`	`75`
`76`	`76`	`struct WeakIValueEq {`
`77`	`77`	`bool operator()(const WeakIValue& t1, const WeakIValue& t2) const {`
`78`		`- return t1.lock().isSameIdentity(t2.lock());`
	`78`	`+ return t1.isSameIdentity(t2);`
`79`	`79`	`}`
`80`	`80`	`};`
`81`	`81`