Simplified the QLoss layer