Tensorflow 2 Semantic Segmentation - loss function for two classes

I am trying to implement a U-net for semantic segmentation with two classes (foreground=1 and background=0) in the segmentation mask images, following this tutorial.

They have used SparseCategoricalCrossentropy for OUTPUT_CLASSES = 3, as shown below:

OUTPUT_CLASSES = 3

model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam',
          loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
          metrics=['accuracy'])

EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_batches, epochs=EPOCHS,
                      steps_per_epoch=STEPS_PER_EPOCH,
                      validation_steps=VALIDATION_STEPS,
                      validation_data=test_batches,
                      callbacks=[DisplayCallback()])

If I use the same network with OUTPUT_CLASSES = 2 the training loss is NaN.

What should I use for the loss function since I only have 2 classes unlike the example dataset (Oxford IIIT Pets) used here which has 3 classes.