With num_workers=0, I get:
Epoch: 0
Batch 0
Batch 1
Batch 2
Traceback (most recent call last):
  File "sv_unet/main.py", line 246, in <module>
    training(unet, train_loader, valid_loader, epochs=epochs, batch_size=BATCH_SIZE, device=device, fold_dir=fold_dir, restore=RESTORE)
  File "/content/drive/My Drive/TFM/monai_unet/sv_unet/training.py", line 119, in training
    for batch in train_loader:
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py", line 345, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py", line 385, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/_utils/collate.py", line 79, in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/_utils/collate.py", line 79, in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/_utils/collate.py", line 55, in default_collate
    return torch.stack(batch, 0, out=out)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/wrap.py", line 89, in wrapper
    return orig_fn(seq, *args, **kwargs)
RuntimeError: Expected object of scalar type double but got scalar type float for sequence element 1.