Custom Loss for Pretrained TF HF Models

Beginners
1

Maybe a bit of a dumb question but I’m trying to fine tune a simple POS tagger like so:

    model = transformers.TFAutoModelForTokenClassification.from_pretrained('distilbert-base-multilingual-cased', num_labels=len(LABELS))
    optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    model.compile(optimizer=optimizer,
                  loss={"loss": loss},
                  metrics=tf.keras.metrics.SparseCategoricalAccuracy())
    model.fit(train,
              epochs=args.epochs)

but it complains that the target for getting the gradient is None:

TypeError: Target should be a list or nested structure of Tensors or Variables to be differentiated, but recieved None

I tried passing the loss directly and as a dictionary like in the snippet here and neither way works. How do I pass it a custom loss properly?

I’m at a loss since the example code pretty much does the same thing I’m doing.