System Info

transformers/pytorch_utils.py", line 300, in id_tensor_storage
if is_torch_greater_or_equal_than_2_0:
from torch.distributed.tensor import DTensor

The error "ImportError: cannot import name 'dtensor' from 'torch.distributed.tensor'" arises due to changes in the location of the DTensor class within PyTorch's distributed package, specifically after version 2.5. Not after 2.0.
So "if is_torch_greater_or_equal_than_2_0: " should be "if is_torch_greater_or_equal_than_2_5:"

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Traceback (most recent call last):
File "/blue/bphl-florida/dongyibo/PPML/Geneformer/geneformer2/Geneformer/cancer_type_classify.py", line 25, in
all_metrics = cc.validate(model_directory="/blue/bphl-florida/dongyibo/PPML/Geneformer/geneformer2/Geneformer/gf-12L-95M-i4096_CLcancer",
File "/blue/bphl-florida/dongyibo/PPML/Geneformer/geneformer2/Geneformer/geneformer/classifier.py", line 794, in validate
trainer = self.train_classifier(
File "/blue/bphl-florida/dongyibo/PPML/Geneformer/geneformer2/Geneformer/geneformer/classifier.py", line 1282, in train_classifier
trainer.train()
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 2240, in train
return inner_training_loop(
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 2656, in _inner_training_loop
self._maybe_log_save_evaluate(
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 3102, in _maybe_log_save_evaluate
self._save_checkpoint(model, trial)
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 3199, in _save_checkpoint
self.save_model(output_dir, _internal_call=True)
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 3911, in save_model
self._save(output_dir)
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/trainer.py", line 4015, in _save
self.model.save_pretrained(
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/modeling_utils.py", line 3572, in save_pretrained
ptrs[id_tensor_storage(tensor)].append(name)
File "/blue/bphl-florida/dongyibo/conda/envs/DL/lib/python3.9/site-packages/transformers/pytorch_utils.py", line 300, in id_tensor_storage
from torch.distributed.tensor import DTensor
ImportError: cannot import name 'DTensor' from 'torch.distributed.tensor'

Expected behavior

pytorch_utils.py", line 299, in id_tensor_storage

"if is_torch_greater_or_equal_than_2_0:" should be "if is_torch_greater_or_equal_than_2_5:"