Add PyTorch test file (#720)

lutzroeder · Apr 25, 2022 · 798b35d · 798b35d
1 parent f9beec5
commit 798b35d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/source/pytorch.js b/source/pytorch.js
@@ -693,6 +693,7 @@ pytorch.Tensor = class {
             case 'float16':
             case 'float32':
             case 'float64':
+            case 'bfloat16':
                 break;
             default:
                 context.state = "Tensor data type '" + this._type.dataType + "' is not supported.";
@@ -778,6 +779,11 @@ pytorch.Tensor = class {
                         context.index += 8;
                         context.count++;
                         break;
+                    case 'bfloat16':
+                        results.push(context.dataView.getBfloat16(context.index, this._littleEndian));
+                        context.index += 2;
+                        context.count++;
+                        break;
                     default:
                         throw new pytorch.Error("Unsupported tensor data type '" + context.dataType + "'.");
                 }
@@ -1949,6 +1955,11 @@ pytorch.Execution = class extends python.Execution {
                 super(size, torch.qint32);
             }
         });
+        this.registerType('torch.BFloat16Storage', class extends torch.storage._StorageBase {
+            constructor(size) {
+                super(size, torch.bfloat16);
+            }
+        });
         this.registerType('torch.Size', class extends Array {
             constructor(size) {
                 super(size.length);
@@ -2043,6 +2054,7 @@ pytorch.Execution = class extends python.Execution {
         this.registerType('torch.QInt8Tensor', class extends torch.Tensor {});
         this.registerType('torch.QUInt8Tensor', class extends torch.Tensor {});
         this.registerType('torch.QInt32Tensor', class extends torch.Tensor {});
+        this.registerType('torch.BFloat16Tensor', class extends torch.Tensor {});
         this.registerType('torch.cuda.FloatTensor', class extends torch.Tensor {});
         this.registerType('torch.cuda.DoubleTensor', class extends torch.Tensor {});
         torch.uint8 = new torch.dtype(pytorch.ScalarType.uint8);

diff --git a/test/models.json b/test/models.json
@@ -4600,6 +4600,13 @@
     "format":   "PyTorch v0.1.10",
     "link":     "https://github.com/facebookresearch/kill-the-bits/tree/master/src/models/compressed"
   },
+  {
+    "type":     "pytorch",
+    "target":   "mnist_bfloat16.pt",
+    "source":   "https://github.com/lutzroeder/netron/files/8556279/mnist_bfloat16.pt.zip[mnist_bfloat16.pt]",
+    "format":   "PyTorch v1.6",
+    "link":     "https://github.com/lutzroeder/netron/issues/720"
+  },
   {
     "type":     "pytorch",
     "target":   "mnist_linear.ckpt",