TorchScript 1.4 prototype (#281)

lutzroeder · Jan 19, 2020 · 6074cc9 · 6074cc9
1 parent fb1d08d
commit 6074cc9
Show file tree

Hide file tree

Showing 3 changed files with 252 additions and 179 deletions.
diff --git a/src/pickle.js b/src/pickle.js
@@ -228,7 +228,12 @@ pickle.Unpickler = class {
                     const state = stack.pop();
                     let obj = stack.pop();
                     if (obj.__setstate__) {
-                        obj.__setstate__(state);
+                        if (obj.__setstate__.__call__) {
+                            obj.__setstate__.__call__([ obj, state ]);
+                        }
+                        else {
+                            obj.__setstate__(state);
+                        }
                     }
                     else {
                         for (const p in state) {

diff --git a/src/pytorch-metadata.json b/src/pytorch-metadata.json
@@ -845,7 +845,51 @@
     }
   },
   {
-    "name": "torch.conv2d_relu",
+    "name": "ops.quantized.conv2d",
+    "schema": {
+      "attributes": [
+        {
+          "default": 1,
+          "name": "stride",
+          "type": "int64[]"
+        },
+        {
+          "default": 0,
+          "name": "padding",
+          "type": "int64[]"
+        },
+        {
+          "default": 1,
+          "name": "dilation",
+          "type": "int64[]"
+        },
+        {
+          "default": 1,
+          "name": "groups",
+          "type": "int64"
+        }
+      ],
+      "category": "Layer",
+      "inputs": [
+        {
+          "name": "input"
+        },
+        {
+          "name": "weight"
+        },
+        {
+          "name": "bias"
+        }
+      ],
+      "outputs": [
+        {
+          "name": "output"
+        }
+      ]
+    }
+  },
+  {
+    "name": "ops.quantized.conv2d_relu",
     "schema": {
       "attributes": [
         {
@@ -893,6 +937,12 @@
       ]
     }
   },
+  {
+    "name": "ops.quantized.linear",
+    "schema": {
+      "category": "Layer"
+    }
+  },
   {
     "name": "torch.max_pool2d",
     "schema": {