I am trying to implement a t-SNE visualization in tensorflow for an image classification task. What I mainly found on the net have all been implemented in Pytorch. See

model2 outputs the features you want to visualize and model outputs the predicted classes with the help of np.argmax. Also, this example is using a dataset with 5 classes, that is why there are 5 different colors. In your case, you only have 2 classes and therefore 2 colors.