Weighting and loss function for multi-dimensional output on ECG neural network in Tensorflow

Question

I am working on a DNN that is training on ecg data with a shape of [None,1,2500] and output shape of [None,12,19] where 19 is a vector of 1 second bins and each of the 12 rows is a distinct binary class, 0,1. Classes are NOT exclusive. I am currently using sigmoid activation and binary crossentropy. The goal is to use weighted binary crossentropy as many of the classes are imbalanced in favor of the negative class i.e. 0. Using the off the shelf Tensorflow binarycrossentropy function breaks training with error,

    File "/usr/lib/python3/dist-packages/keras/engine/training.py", line 1051, in train_function  *
        return step_function(self, iterator)
    File "/usr/lib/python3/dist-packages/keras/engine/training.py", line 1040, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/martin/.local/lib/python3.8/site-packages/six.py", line 719, in reraise
        raise value
    File "/usr/lib/python3/dist-packages/keras/engine/training.py", line 1030, in run_step  **
        outputs = model.train_step(data)
    File "/usr/lib/python3/dist-packages/keras/engine/training.py", line 890, in train_step
        loss = self.compute_loss(x, y, y_pred, sample_weight)
    File "/usr/lib/python3/dist-packages/keras/engine/training.py", line 948, in compute_loss
        return self.compiled_loss(
    File "/usr/lib/python3/dist-packages/keras/engine/compile_utils.py", line 201, in __call__
        loss_value = loss_obj(y_t, y_p, sample_weight=sw)
    File "/usr/lib/python3/dist-packages/keras/losses.py", line 140, in __call__
        return losses_utils.compute_weighted_loss(
    File "/usr/lib/python3/dist-packages/keras/utils/losses_utils.py", line 326, in compute_weighted_loss
        losses, _, sample_weight = squeeze_or_expand_dimensions(  # pylint: disable=unbalanced-tuple-unpacking
    File "/usr/lib/python3/dist-packages/keras/utils/losses_utils.py", line 212, in squeeze_or_expand_dimensions
        sample_weight = tf.squeeze(sample_weight, [-1])
    ValueError: Can not squeeze dim[1], expected a dimension of 1, got 19 for '{{node binary_crossentropy/weighted_loss/Squeeze}} = Squeeze[T=DT_FLOAT, squeeze_dims=[-1]](cond/Identity_4)' with input shapes: [?,19].

The failure of applying weights to the output with Tensorflows off-the-shelf binarycrossentropy spurred the use of a custom loss function as seen below, however, performance of the model with this custom loss is weak.

class MultiLabelLoss(tf.keras.losses.Loss):
    def __init__(self, pos_weights, neg_weights, **kwargs):
        self.pos_weights = pos_weights
        self.neg_weights = neg_weights
        super().__init__(**kwargs)
    def call(self, y_true, y_logit):
        """
        Multi-label cross-entropy
        * Required "Wp", "Wn" as positive & negative class-weights
        y_true: true value
        y_logit: predicted value
        """
        arrh_classes = ArrhClassType
        class_count = ArrhClassType.get_active_class_count()
        loss = float(0)
        # wp = self.class_weights['positive_weights']
        # wn = self.class_weights['negative_weights']
        wp = self.pos_weights
        wn = self.neg_weights
        #print("LOGITS: ", y_logit.numpy().shape)
        for i in range(class_count):
            class_key = arrh_classes.find_from_annot_id(i).value['short']
            if class_key not in wp:
                wp[class_key] = 0.5
                wn[class_key] = 0.5
            log_logit = K.log(y_logit[:,:,i] + K.epsilon())
            true_val = y_true[:,:,i]
            ann = wp[class_key]
            first_term = ann * true_val * log_logit
            second_term = wn[class_key] * (1 - y_true[:,:,i]) * K.log(1 - y_logit[:,:,i] + K.epsilon())
            loss -= (first_term + second_term)
        return loss
    def get_config(self):
        base_config = super().get_config()
        self.pos_weights = {key:value.numpy() if type(value)==tf_python.framework.ops.EagerTensor else value for key,value in self.pos_weights.items()}
        self.neg_weights = {key:value.numpy() if type(value)==tf_python.framework.ops.EagerTensor else value for key,value in self.neg_weights.items()}

        return {**base_config, "pos_weights": self.pos_weights, "neg_weights": self.neg_weights}

So the question is, is a custom loss function necessary to train my model with weighting with the output structure I have defined? If not, how can I apply weighting to Tensorflows binarycrossentropy? If yes, how can I define an effective custom binarycrossentropy loss function for my classification problem with weighting for my imbalanced classes.

Finally, is my approach completely incorrect? Should I be structuring my outputs in a completely different way? Any and all insight is welcome and appreciated.

Thank you