Set pathOrIOHandler to IOHandler

infinitered · Aug 1, 2019 · 6841bb7 · 6841bb7
1 parent 697b0ad
commit 6841bb7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/index.ts b/src/index.ts
@@ -43,7 +43,7 @@ export class NSFWJS {
   public endpoints: string[]
 
   private options: nsfwjsOptions
-  private pathOrIOHandler: string
+  private pathOrIOHandler: string | IOHandler
   private model: tf.LayersModel
   private intermediateModels: { [layerName: string]: tf.LayersModel } = {}
 
@@ -59,6 +59,7 @@ export class NSFWJS {
     if (typeof modelPathBaseOrIOHandler === 'string') {
       this.pathOrIOHandler = `${modelPathBaseOrIOHandler}model.json`
     } else {
+      this.pathOrIOHandler = modelPathBaseOrIOHandler
     }
   }