# Buggy code
def update_metrics(self, logits, labels, tagmask=None):
    logits, labels = logits.detach(), labels.detach()

    def make_one_hot(batch, depth=2):
        """
        Creates a one-hot embedding of dimension 2.
        Parameters:
        batch: list of size batch_size of class predictions
        Returns:
        one hot encoding of size [batch_size, 2]
        """
        ones = torch.sparse.torch.eye(depth).cuda()
        return ones.index_select(0, batch)

    binary_preds = make_one_hot(logits, depth=2)
    # Make label_ints a batch_size list of labels
    label_ints = torch.argmax(labels, dim=1)
    self.f1_scorer(binary_preds, label_ints)
    self.acc_scorer(binary_preds.long(), labels.long())


# Traceback (ValueError)
Traceback (most recent call last):
File "main.py", line 566, in <module>
main(sys.argv[1:])
File "main.py", line 534, in main
phase="target_train",
File "/Users/srbowman/jiant/src/trainer.py", line 594, in train
output_dict = self._forward(batch, task=task)
File "/Users/srbowman/jiant/src/trainer.py", line 1035, in _forward
model_out = self._model.forward(task, tensor_batch)
File "/Users/srbowman/jiant/src/models.py", line 830, in forward
out = self._span_forward(batch, task, predict)
File "/Users/srbowman/jiant/src/models.py", line 917, in _span_forward
out = module.forward(batch, sent_embs, sent_mask, task, predict)
File "/Users/srbowman/jiant/src/modules/span_modules.py", line 134, in forward
out["loss"] = self.compute_loss(logits, batch["labels"].squeeze(dim=1), task)
File "/Users/srbowman/jiant/src/modules/span_modules.py", line 184, in compute_loss
raise ValueError("Unsupported loss type '%s' " "for edge probing." % self.loss_type)
ValueError: Unsupported loss type '' for edge probing.
Traceback (most recent call last):
File "main.py", line 577, in <module>
raise e  # re-raise exception, in case debugger is attached.
File "main.py", line 566, in <module>
main(sys.argv[1:])
File "main.py", line 534, in main
phase="target_train",
File "/Users/srbowman/jiant/src/trainer.py", line 594, in train
output_dict = self._forward(batch, task=task)
File "/Users/srbowman/jiant/src/trainer.py", line 1035, in _forward
model_out = self._model.forward(task, tensor_batch)
File "/Users/srbowman/jiant/src/models.py", line 830, in forward
out = self._span_forward(batch, task, predict)
File "/Users/srbowman/jiant/src/models.py", line 917, in _span_forward
out = module.forward(batch, sent_embs, sent_mask, task, predict)
File "/Users/srbowman/jiant/src/modules/span_modules.py", line 134, in forward
out["loss"] = self.compute_loss(logits, batch["labels"].squeeze(dim=1), task)
File "/Users/srbowman/jiant/src/modules/span_modules.py", line 184, in compute_loss
raise ValueError("Unsupported loss type '%s' " "for edge probing." % self.loss_type)
ValueError: Unsupported loss type '' for edge probing.

# Fixed code
def update_metrics(self, logits, labels, tagmask=None):
    logits, labels = logits.detach(), labels.detach()

    def make_one_hot(batch, depth=2):
        """
        Creates a one-hot embedding of dimension 2.
        Parameters:
        batch: list of size batch_size of class predictions
        Returns:
        one hot encoding of size [batch_size, 2]
        """
        ones = torch.sparse.torch.eye(depth)
        if torch.cuda.is_available():
            ones = ones.cuda()
        return ones.index_select(0, batch)

    binary_preds = make_one_hot(logits, depth=2)
    # Make label_ints a batch_size list of labels
    label_ints = torch.argmax(labels, dim=1)
    self.f1_scorer(binary_preds, label_ints)
    self.acc_scorer(binary_preds.long(), labels.long())