OK so I\'m trying to implement this custom loss I read in a paper and it\'s my first custom loss so I might be missing something obvious, but I get the error shown below. Fr