Skip to content

Commit

Permalink
redirect to cpu only for webgl backend (tensorflow#538)
Browse files Browse the repository at this point in the history
  • Loading branch information
pyu10055 authored Nov 24, 2020
1 parent 4a538cb commit 7c2b3e4
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions coco-ssd/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,9 @@ export class ObjectDetection {

const prevBackend = tf.getBackend();
// run post process in cpu
tf.setBackend('cpu');
if (tf.getBackend() === 'webgl') {
tf.setBackend('cpu');
}
const indexTensor = tf.tidy(() => {
const boxes2 =
tf.tensor2d(boxes, [result[1].shape[1], result[1].shape[3]]);
Expand All @@ -151,7 +153,9 @@ export class ObjectDetection {
indexTensor.dispose();

// restore previous backend
tf.setBackend(prevBackend);
if (prevBackend !== tf.getBackend()) {
tf.setBackend(prevBackend);
}

return this.buildDetectedObjects(
width, height, boxes, maxScores, indexes, classes);
Expand Down

0 comments on commit 7c2b3e4

Please sign in to comment.