update try/catch

Author: alicialics

tensorboard/plugins/projector/vz_projector/knn.ts

@@ -49,16 +49,14 @@ export function findKNNGPUCosDistNorm<T>(
   // pair of points, which we sort using KMin data structure to obtain the
   // K nearest neighbors for each point.
   const nearest: NearestEntry[][] = new Array(N);
-
-  const typedArray = vector.toTypedArray(dataPoints, accessor);
-  const bigMatrix = tf.tensor(typedArray, [N, dim]);
-  const bigMatrixTransposed = tf.transpose(bigMatrix);
-
   function step(resolve: (result: NearestEntry[][]) => void) {
     util
       .runAsyncTask(
         'Finding nearest neighbors...',
         async () => {
+          const typedArray = vector.toTypedArray(dataPoints, accessor);
+          const bigMatrix = tf.tensor(typedArray, [N, dim]);
+          const bigMatrixTransposed = tf.transpose(bigMatrix);
           // 1 - A * A^T.
           const bigMatrixSquared = tf.matMul(bigMatrix, bigMatrixTransposed);
           const cosDistMatrix = tf.sub(1, bigMatrixSquared);
@@ -68,6 +66,9 @@ export function findKNNGPUCosDistNorm<T>(
           // [ 3  4 ],
           // `.data()` returns [1, 2, 3, 4].
           const partial = await cosDistMatrix.data();
+          // Discard all tensors and free up the memory.
+          bigMatrix.dispose();
+          bigMatrixTransposed.dispose();
           bigMatrixSquared.dispose();
           cosDistMatrix.dispose();
           for (let i = 0; i < N; i++) {
@@ -93,15 +94,9 @@ export function findKNNGPUCosDistNorm<T>(
       .then(
         () => {
           logging.setModalMessage(null!, KNN_MSG_ID);
-          // Discard all tensors and free up the memory.
-          bigMatrix.dispose();
-          bigMatrixTransposed.dispose();
           resolve(nearest);
         },
         (error) => {
-          // Discard all tensors and free up the memory.
-          bigMatrix.dispose();
-          bigMatrixTransposed.dispose();
           // GPU failed. Reverting back to CPU.
           logging.setModalMessage(null!, KNN_MSG_ID);
           let distFunc = (a, b, limit) => vector.cosDistNorm(a, b);