Update pytorch_lightning/plugins/precision/double.py

ethanwharris · justusschock · web-flow · commit 4816bda0dd93 · 2021-03-23T08:52:38.000Z
Co-authored-by: Justus Schock &lt;12886177+justusschock@users.noreply.github.com&gt;
diff --git a/pytorch_lightning/plugins/precision/double.py b/pytorch_lightning/plugins/precision/double.py
@@ -60,7 +60,7 @@ def new_method(*args, **kwargs) -> Any:
             )
 
         setattr(model, method_name, new_method if callable(old_method) else old_method)
-        return _DoublePrecisionPatch(model, method_name, old_method)
+        return cls(model, method_name, old_method)
 
 
 class DoublePrecisionPlugin(PrecisionPlugin):

Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ def new_method(args, *kwargs) -> Any:`
`60`	`60`	`)`
`61`	`61`
`62`	`62`	`setattr(model, method_name, new_method if callable(old_method) else old_method)`
`63`		`- return _DoublePrecisionPatch(model, method_name, old_method)`
	`63`	`+ return cls(model, method_name, old_method)`
`64`	`64`
`65`	`65`
`66`	`66`	`class DoublePrecisionPlugin(PrecisionPlugin):`