Categoryencoding (#51)

* added reviewer requests
CodeSmithDSMLProjects · Oct 4, 2022 · 1eb6f6c · 1eb6f6c
1 parent 7222f30
commit 1eb6f6c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tfjs-layers/src/layers/preprocessing/category_encoding.ts b/tfjs-layers/src/layers/preprocessing/category_encoding.ts
@@ -85,9 +85,9 @@ export class CategoryEncoding extends Layer {
               `countWeights is not used when outputMode !== count.
               Received countWeights=${kwargs['countWeights']}`);
           }
-          if (countWeightsArg.rank === 1 || countWeightsArg.rank === 2) {
-            const countWeights = getExactlyOneTensor(kwargs['countWeights']) as Tensor1D | Tensor2D;
-          }
+
+          countWeights
+            =  getExactlyOneTensor(kwargs['countWeights']) as Tensor1D|Tensor2D;
         }
 
         const maxValue = max(inputs);