Help with Entity Embeddings in Pytorch

I tried to build Entity Embeddings of categorical data of a dataset.
I took a dataset - "Bike share”.This dataset shows number of bike share/rent/sales in every month.The month column has 12 categories.

I wanted to make 12 Entity embeddings or 12 vectors, one for each month.I trained the embeddings against the bike share/rent/sales or i should say i used bike share/rent/sales as my labels.To visualise the embeddings I plotted the learned Embeddings/vectors
On a 3D graph and I expected that the months with similar sales will be close together on the 3D graph but I didn’t get the expected results.
In fact when I made the same model in Keras I got the results I was expecting.For example the batch size of 8 works very well in keras but same batch size increases loss in pytorch. In pytorch big batch size(40) reduces loss.

Would anyone be kind enough to have a look at my code in pytorch.It’s a very short code and will only take 10 minutes.

Here’s the link to my Pytorch code - https://github.com/akshay6893/Entity_Embeddings/blob/master/Entity%20Embedding%20in%20Pytorch%202.ipynb

and keras code - https://github.com/akshay6893/Entity_Embeddings/blob/master/Entity_Embeddings_keras.ipynb

Here’s the data online - https://www.kaggle.com/marklvl/bike-sharing-dataset.(very small dataset)