This PR fixes https://github.com/tensorflow/tfjs/issues/1329 

This PR fixes https://github.com/tensorflow/tfjs/issues/1342