[Retiarii] fix missed "attributes" in ToDevice pytorch op (#4266)

7fe57a0e · Zhenhua Han · GitHub · 99c0f938 · 7fe57a0e
Unverified Commit 7fe57a0e authored Oct 22, 2021 by Zhenhua Han Committed by GitHub Oct 22, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

nni/retiarii/operation_def/torch_op_def.py nni/retiarii/operation_def/torch_op_def.py +2 -1

No files found.
--- a/nni/retiarii/operation_def/torch_op_def.py
+++ b/nni/retiarii/operation_def/torch_op_def.py
@@ -491,7 +491,8 @@ class AtenAvgpool2d(PyTorchOperation):
 class ToDevice(PyTorchOperation):
    _artificial_op_name = "ToDevice"

-    def __init__(self, type_name: str, parameters: Dict[str, Any], _internal: bool = False):
+    def __init__(self, type_name: str, parameters: Dict[str, Any], _internal: bool = False,
+                 attributes: Dict[str, Any] = None):
        self.type = "ToDevice"
        self.device = parameters['device']
        self.overridden_device_repr = None