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

gagan3012: <p>Im using torch 2.0.1 and transfomers 4.34.0dev</p>
<p>I keep getting this error:</p>
<blockquote>
<p>Blockquote<br>
Hello,<br>
I fixed all broken packages and I still get the same error even on GPU.<br>
[gagan30@ng20603 diff]$ cd …<br>
[gagan30@ng20603 arocr]$ cd …<br>
[gagan30@ng20603 scratch]$ cd …<br>
[gagan30@ng20603 ~]$ cd scratch/arocr/diff/<br>
[gagan30@ng20603 diff]$ ./train.sh</p>
</blockquote>
<p>Lmod is automatically replacing “intel/2020.1.217” with “gcc/9.3.0”.</p>
<p>Due to MODULEPATH changes, the following have been reloaded:</p>
<ol>
<li>blis/0.8.1     2) flexiblas/3.0.4     3) openmpi/4.0.3     4) ucx/1.8.0</li>
</ol>
<p>The following have been reloaded with a version change:</p>
<ol>
<li>libfabric/1.10.1 =&gt; libfabric/1.15.1</li>
</ol>
<p>[2023-09-07 15:52:27,968] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)</p>
<blockquote>
<p>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=&lt;function run_dreambooth_command_factory at 0x1550a2a4b130&gt;)<br>
INFO    Running DreamBooth Training<br>
WARNING Parameters supplied but not used: version, func<br>
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.<br>
ERROR   train has failed due to an exception:<br>
ERROR   Traceback (most recent call last):<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/utils.py”, line 280, in wrapper<br>
return func(*args, **kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/trainers/dreambooth/<strong>main</strong>.py”, line 103, in train<br>
tokenizers, text_encoders, vae, unet, noise_scheduler = utils.load_model_components(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/autotrain/trainers/dreambooth/utils.py”, line 263, in load_model_components<br>
vae = AutoencoderKL.from_pretrained(config.model, subfolder=“vae”, revision=config.revision)</p>
</blockquote>
<p>File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/modeling_utils.py”, line 611, in from_pretrained<br>
model = cls.from_config(config, **unused_kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 254, in from_config<br>
model = cls(**init_dict)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 636, in inner_init<br>
init(self, *args, **init_kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/autoencoder_kl.py”, line 95, in <strong>init</strong><br>
self.encoder = Encoder(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/vae.py”, line 55, in <strong>init</strong><br>
self.conv_in = torch.nn.Conv2d(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 450, in <strong>init</strong><br>
super().<strong>init</strong>(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 144, in <strong>init</strong><br>
self.reset_parameters()<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 150, in reset_parameters<br>
init.kaiming_uniform_(self.weight, a=math.sqrt(5))<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/init.py”, line 412, in kaiming_uniform_<br>
return tensor.uniform_(-bound, bound)<br>
File "/home/gagan30/ENV/lib/python3.10/site-packages/torch/<em>decomp/decompositions.py", line 1958, in uniform</em><br>
return self.copy_((high - low) * torch.rand_like(self) + low)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 220, in _fn<br>
result = fn(*args, **kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 130, in _fn<br>
result = fn(**bound.arguments)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/<strong>init</strong>.py”, line 926, in _ref<br>
return prim(a, b)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/<strong>init</strong>.py”, line 1532, in mul<br>
return prims.mul(a, b)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_ops.py”, line 287, in <strong>call</strong><br>
return self._op(*args, **kwargs or {})<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims/<strong>init</strong>.py”, line 350, in <em>elementwise_meta<br>
shape = utils.extract_shape(*args</em>, allow_cpu_scalar_tensors=True)<br>
AttributeError: module ‘torch._prims.utils’ has no attribute ‘extract_shape’</p>
<p>[2023-09-07 15:52:35,501] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)<br>
Loading pipeline components…:  14%|███████▍                                            | 1/7 [00:00&lt;00:01,  3.79it/s]<br>
Traceback (most recent call last):<br>
File “/lustre07/scratch/gagan30/arocr/diff/gen.py”, line 82, in <br>
Fire(generate_images)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 141, in Fire<br>
component_trace = <em>Fire(component, args, parsed_flag_args, context, name)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 466, in <em>Fire<br>
component, remaining_args = <em>CallAndUpdateTrace(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/fire/core.py”, line 681, in <em>CallAndUpdateTrace<br>
component = fn(*varargs, **kwargs)<br>
File “/lustre07/scratch/gagan30/arocr/diff/gen.py”, line 41, in generate_images<br>
pipe = DiffusionPipeline.from_pretrained(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py”, line 1093, in from_pretrained<br>
loaded_sub_model = load_sub_model(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py”, line 467, in load_sub_model<br>
loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/modeling_utils.py”, line 611, in from_pretrained<br>
model = cls.from_config(config, **unused_kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 254, in from_config<br>
model = cls(**init_dict)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/configuration_utils.py”, line 636, in inner_init<br>
init(self, *args, **init_kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/diffusers/models/unet_2d_condition.py”, line 263, in <strong>init</strong><br>
self.conv_in = nn.Conv2d(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 450, in <strong>init</strong><br>
super().<strong>init</strong>(<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 144, in <strong>init</strong><br>
self.reset_parameters()<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/modules/conv.py”, line 150, in reset_parameters<br>
init.kaiming_uniform</em>(self.weight, a=math.sqrt(5))<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/nn/init.py”, line 412, in kaiming_uniform</em><br>
return tensor.uniform</em>(-bound, bound)<br>
File "/home/gagan30/ENV/lib/python3.10/site-packages/torch/<em>decomp/decompositions.py", line 1958, in uniform</em><br>
return self.copy</em>((high - low) * torch.rand_like(self) + low)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 220, in _fn<br>
result = fn(*args, **kwargs)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims_common/wrappers.py”, line 130, in _fn<br>
result = fn(**bound.arguments)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/<strong>init</strong>.py”, line 926, in _ref<br>
return prim(a, b)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_refs/<strong>init</strong>.py”, line 1532, in mul<br>
return prims.mul(a, b)<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_ops.py”, line 287, in <strong>call</strong><br>
return self._op(*args, **kwargs or {})<br>
File “/home/gagan30/ENV/lib/python3.10/site-packages/torch/_prims/<strong>init</strong>.py”, line 350, in <em>elementwise_meta<br>
shape = utils.extract_shape(*args</em>, allow_cpu_scalar_tensors=True)<br>
AttributeError: module ‘torch._prims.utils’ has no attribute ‘extract_shape’<br>
[gagan30@ng20603 diff]$</p>

ptrblck: <p>I can properly import the function in <code>2.0.1</code> and the latest nightly:</p>
<pre><code class="lang-python">torch.__version__
&gt; '2.0.1+cu117'
import torch._prims
torch._prims.utils.extract_shape
&gt; &lt;function extract_shape at 0x7fc6303917e0&gt;
</code></pre>
<pre><code class="lang-auto">torch.__version__
&gt; '2.2.0.dev20230907+cu121'
import torch._prims
torch._prims.utils.extract_shape
&gt; &lt;function torch._prims_common.extract_shape(*args, allow_cpu_scalar_tensors: 'bool') -&gt; 'Optional[ShapeType]'&gt;
</code></pre>