Add PyTorch PReLU (#133)

lutzroeder · Oct 26, 2018 · d5e326a · d5e326a
1 parent 8cb1a10
commit d5e326a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/src/pytorch-metadata.json b/src/pytorch-metadata.json
@@ -140,6 +140,19 @@
       "package": "torch.nn.modules.activation"
     }
   },
+  {
+    "name": "PReLU",
+    "schema": {
+      "attributes": [
+        {
+          "name": "training",
+          "visible": false
+        }
+      ],
+      "category": "Activation",
+      "package": "torch.nn.modules.activation"
+    }
+  },
   {
     "name": "MaxPool2d",
     "schema": {

diff --git a/src/pytorch-model.js b/src/pytorch-model.js
@@ -73,6 +73,7 @@ class PyTorchModelFactory {
 
             constructorTable['argparse.Namespace'] = function (args) { this.args = args; };
             constructorTable['torch.nn.modules.activation.ReLU'] = function () {};
+            constructorTable['torch.nn.modules.activation.PReLU'] = function () {};
             constructorTable['torch.nn.modules.activation.Tanh'] = function () {};
             constructorTable['torch.nn.modules.activation.Sigmoid'] = function () {};
             constructorTable['torch.nn.modules.batchnorm.BatchNorm2d'] = function () {};