AttributeError: module 'torch._prims.utils' has no attribute 'extract_shape'

Im using torch 2.0.1 and transfomers 4.34.0dev

I keep getting this error:

Blockquote
Hello,
I fixed all broken packages and I still get the same error even on GPU.
[gagan30@ng20603 diff]$ cd …
[gagan30@ng20603 arocr]$ cd …
[gagan30@ng20603 scratch]$ cd …
[gagan30@ng20603 ~]$ cd scratch/arocr/diff/
[gagan30@ng20603 diff]$ ./train.sh

Lmod is automatically replacing “intel/2020.1.217” with “gcc/9.3.0”.

Due to MODULEPATH changes, the following have been reloaded:

  1. blis/0.8.1 2) flexiblas/3.0.4 3) openmpi/4.0.3 4) ucx/1.8.0

The following have been reloaded with a version change:

  1. libfabric/1.10.1 => libfabric/1.15.1

[2023-09-07 15:52:27,968] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)

INFO Namespace(version=False, model=‘stable-diffusion-xl-base-1.0’, revision=None, tokenizer=None, image_path=‘images/gagan/’, class_image_path=None, prompt=‘photo of gagan’, class_prompt=None, num_class_images=100, class_labels_conditioning=None, prior_preservation=None, prior_loss_weight=1.0, project_name=‘models/gagan’, seed=42, resolution=1024, center_crop=None, train_text_encoder=None, batch_size=1, sample_batch_size=4, epochs=1, num_steps=1000, checkpointing_steps=100000, resume_from_checkpoint=None, gradient_accumulation=4, gradient_checkpointing=True, lr=0.0001, scale_lr=None, scheduler=‘constant’, warmup_steps=0, num_cycles=1, lr_power=1.0, dataloader_num_workers=0, use_8bit_adam=True, adam_beta1=0.9, adam_beta2=0.999, adam_weight_decay=0.01, adam_epsilon=1e-08, max_grad_norm=1.0, allow_tf32=None, prior_generation_precision=None, local_rank=-1, xformers=None, pre_compute_text_embeddings=None, tokenizer_max_length=None, text_encoder_use_attention_mask=None, rank=4, xl=None, fp16=True, bf16=None, token=None, repo_id=None, push_to_hub=None, validation_prompt=None, num_validation_images=4, validation_epochs=50, checkpoints_total_limit=None, validation_images=None, logging=None, username=None, func=<function run_dreambooth_command_factory at 0x1550a2a4b130>)
INFO Running DreamBooth Training
WARNING Parameters supplied but not used: version, func
You are using a model of type clip_text_model to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
ERROR train has failed due to an exception:
ERROR Traceback (most recent call last):
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/utils.py”, line 280, in wrapper
return func(*args, **kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/trainers/dreambooth/main.py”, line 103, in train
tokenizers, text_encoders, vae, unet, noise_scheduler = utils.load_model_components(
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/trainers/dreambooth/utils.py”, line 263, in load_model_components
vae = AutoencoderKL.from_pretrained(config.model, subfolder=“vae”, revision=config.revision)

File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/modeling_utils.py”, line 611, in from_pretrained
model = cls.from_config(config, **unused_kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 254, in from_config
model = cls(**init_dict)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 636, in inner_init
init(self, *args, **init_kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/autoencoder_kl.py”, line 95, in init
self.encoder = Encoder(
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/vae.py”, line 55, in init
self.conv_in = torch.nn.Conv2d(
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 450, in init
super().init(
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 144, in init
self.reset_parameters()
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 150, in reset_parameters
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/init.py”, line 412, in kaiming_uniform_
return tensor.uniform_(-bound, bound)
File "/home/gagan30/ENV/lib/python3.10/site-packages/torch/decomp/decompositions.py", line 1958, in uniform
return self.copy_((high - low) * torch.rand_like(self) + low)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 220, in _fn
result = fn(*args, **kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 130, in _fn
result = fn(**bound.arguments)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/init.py”, line 926, in _ref
return prim(a, b)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/init.py”, line 1532, in mul
return prims.mul(a, b)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_ops.py”, line 287, in call
return self._op(*args, **kwargs or {})
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims/init.py”, line 350, in elementwise_meta
shape = utils.extract_shape(*args
, allow_cpu_scalar_tensors=True)
AttributeError: module ‘torch._prims.utils’ has no attribute ‘extract_shape’

[2023-09-07 15:52:35,501] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Loading pipeline components…: 14%|███████▍ | 1/7 [00:00<00:01, 3.79it/s]
Traceback (most recent call last):
File “/lustre07/scratch/gagan30/arocr/diff/gen.py”, line 82, in
Fire(generate_images)
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 141, in Fire
component_trace = Fire(component, args, parsed_flag_args, context, name)
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 466, in Fire
component, remaining_args = CallAndUpdateTrace(
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 681, in CallAndUpdateTrace
component = fn(*varargs, **kwargs)
File “/lustre07/scratch/gagan30/arocr/diff/gen.py”, line 41, in generate_images
pipe = DiffusionPipeline.from_pretrained(
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py”, line 1093, in from_pretrained
loaded_sub_model = load_sub_model(
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py”, line 467, in load_sub_model
loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/modeling_utils.py”, line 611, in from_pretrained
model = cls.from_config(config, **unused_kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 254, in from_config
model = cls(**init_dict)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 636, in inner_init
init(self, *args, **init_kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/unet_2d_condition.py”, line 263, in init
self.conv_in = nn.Conv2d(
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 450, in init
super().init(
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 144, in init
self.reset_parameters()
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 150, in reset_parameters
init.kaiming_uniform
(self.weight, a=math.sqrt(5))
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/init.py”, line 412, in kaiming_uniform

return tensor.uniform
(-bound, bound)
File "/home/gagan30/ENV/lib/python3.10/site-packages/torch/decomp/decompositions.py", line 1958, in uniform
return self.copy
((high - low) * torch.rand_like(self) + low)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 220, in _fn
result = fn(*args, **kwargs)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 130, in _fn
result = fn(**bound.arguments)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/init.py”, line 926, in _ref
return prim(a, b)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/init.py”, line 1532, in mul
return prims.mul(a, b)
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_ops.py”, line 287, in call
return self._op(*args, **kwargs or {})
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims/init.py”, line 350, in elementwise_meta
shape = utils.extract_shape(*args
, allow_cpu_scalar_tensors=True)
AttributeError: module ‘torch._prims.utils’ has no attribute ‘extract_shape’
[gagan30@ng20603 diff]$

I can properly import the function in 2.0.1 and the latest nightly:

torch.__version__
> '2.0.1+cu117'
import torch._prims
torch._prims.utils.extract_shape
> <function extract_shape at 0x7fc6303917e0>
torch.__version__
> '2.2.0.dev20230907+cu121'
import torch._prims
torch._prims.utils.extract_shape
> <function torch._prims_common.extract_shape(*args, allow_cpu_scalar_tensors: 'bool') -> 'Optional[ShapeType]'>