Currently, we do support passing additional arguments in custom loss functions, however we still don't support passing more complicated types, such as vector/matrices, we should do that!