[tfjs-layers[ Fix Layer mem leak for 1D targets

tfjs-layers/src/engine/training.ts

@@ -1535,6 +1535,8 @@ export class LayersModel extends Container implements tfc.InferenceModel {
       lossValues.push(v[0]);
     }
     tfc.dispose(losses);
+    disposeNewTensors(standardizeOut[0], x);
+    disposeNewTensors(standardizeOut[1], y);
     return singletonOrArray(lossValues);
   }
 

tfjs-layers/src/engine/training_tensors.ts

@@ -423,6 +423,8 @@ export async function fitTensors(
   model.isTraining = true;
   let inputs: Tensor[];
   let targets: Tensor[];
+  let originalInputs: Tensor[];
+  let originalTargets: Tensor[];
   let inputValX: Tensor|Tensor[];
   let inputValY: Tensor|Tensor[];
   let valX: Tensor|Tensor[];
@@ -481,8 +483,10 @@ export async function fitTensors(
           Math.floor(inputs[0].shape[0] * (1 - args.validationSplit));
       const originalBatchSize = inputs[0].shape[0];
       valX = sliceArrays(inputs, splitAt, originalBatchSize) as Tensor[];
+      originalInputs = inputs;
       inputs = sliceArrays(inputs, 0, splitAt) as Tensor[];
       valY = sliceArrays(targets, splitAt, originalBatchSize) as Tensor[];
+      originalTargets = targets;
       targets = sliceArrays(targets, 0, splitAt) as Tensor[];
       // TODO(cais): Once sampleWeights becomes available, slice it to get
       //   valSampleWeights.
@@ -537,6 +541,8 @@ export async function fitTensors(
     // Memory clean up.
     disposeNewTensors(inputs, x);
     disposeNewTensors(targets, y);
+    disposeNewTensors(originalInputs, x);
+    disposeNewTensors(originalTargets, y);
     disposeNewTensors(valX as Tensor[], inputValX);
     disposeNewTensors(valY as Tensor[], inputValY);
     if (sampleWeights != null) {

tfjs-layers/src/engine/training_test.ts

@@ -1999,6 +1999,33 @@ describeMathCPUAndWebGL2('LayersModel.fit: No memory leak', () => {
        done();
      });
 
+  it('Repeated fit calls of 1d target leads to no memory leak: validationSplit',
+     async done => {
+       createDenseModelAndData();
+
+       const validationSplit = 0.4;
+       targets = ones([numSamples]);
+       model.compile({optimizer: 'SGD', loss: 'meanSquaredError'});
+       const numTensors0 = memory().numTensors;
+       // Use batchSize === numSamples to get exactly one batch.
+       await model.fit(
+           inputs, targets,
+           {batchSize: 2, epochs: 10, validationSplit, shuffle: true});
+       for (let i = 0; i < 2; ++i) {
+         await model.fit(
+             inputs, targets,
+             {batchSize: 2, epochs: 10, validationSplit, shuffle: true});
+         const numTensorsNow = memory().numTensors;
+         if (numTensorsNow > numTensors0) {
+           done.fail(
+               `Memory leak detected during fit(): Leaked ` +
+               `${numTensorsNow - numTensors0} tensor(s) after the ` +
+               `${i + 1}-th fit() call.`);
+         }
+       }
+       done();
+     });
+
   it('Repeated fit calls leads to no memory leak: validationData',
      async done => {
        createDenseModelAndData();

tfjs-layers/src/utils/test_utils.ts

@@ -95,7 +95,10 @@ export function describeMathCPUAndGPU(testName: string, tests: () => void) {
  */
 export function describeMathCPUAndWebGL2(testName: string, tests: () => void) {
   describeWithFlags(
-      testName, {predicate: testEnv => testEnv.flags['WEBGL_VERSION'] !== 1},
+      testName, {
+        predicate: testEnv =>
+            (testEnv.flags == null || testEnv.flags['WEBGL_VERSION'] === 2)
+      },
       () => {
         tests();
       });
@@ -134,7 +137,8 @@ export function describeMathWebGL2(testName: string, tests: () => void) {
   describeWithFlags(
       testName, {
         predicate: testEnv => testEnv.backendName === 'webgl' &&
-            testEnv.flags['WEBGL_VERSION'] !== 1
+            (testEnv.flags == null || testEnv.flags['WEBGL_VERSION'] === 2)
+
       },
       () => {
         tests();