Hi, one thing I notice is that it seems your learning rate is very high (especially for the Adam optimizer) in comparison to your batch size of 2. Such a high learning rate vs relatively tiny batch size might lead to very noisy learning. I’d suggest reducing the learning rate to 0.001 (basic setting for Adam) and increase the batch size. What are the dimensions of your dataset?
Other than that, I’d suggest to try to strip away complexity until you have a very basic model that works and add steps thereafter to debug your code.