fix TorchModel mistakenly assigning InArgs

XpressAI · Oct 7, 2024 · 8a49889 · 8a49889
1 parent 20d2d16
commit 8a49889
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_nn_layers.py b/torch_nn_layers.py
@@ -20,7 +20,7 @@ class TorchModel(Component):
 
  model_in: InArg[list]
  loss_in: InArg[str]
- learning_rate = InArg[float]
+ learning_rate: InArg[float]
  optimizer_in: InArg[str]
  should_flatten: InArg[bool]
  model_config: OutArg[nn.Module]
@@ -357,4 +357,4 @@ def execute(self, ctx) -> None:
  if self.model_in.value is None:
  self.model_out.value = [nn.Dropout(prob)]
  else:
- self.model_out.value = self.model_in.value + [nn.Dropout(prob)]
+ self.model_out.value = self.model_in.value + [nn.Dropout(prob)]