Tensorflow 2 compatibility

This package does not work with tensorflow 2!