Segmentation fault when running IPEX bf16 example with torch.autocast

y199387: <p>Hi, I am getting a segmentation fault when running <a href="https://intel.github.io/intel-extension-for-pytorch/latest/tutorials/examples.html#complete-bfloat16" rel="noopener nofollow ugc">IPEX BF16 example</a> with <code>torch.autocast</code>.<br>
Environment: pytorch 1.11.0 + intel-extension-for-pytorch 1.11.0<br>
Code to reproduce the error:</p>
<pre><code class="lang-python">import torch
import torchvision
import intel_extension_for_pytorch as ipex

LR = 0.001
DOWNLOAD = True
DATA = 'datasets/cifar10/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(
        root=DATA,
        train=True,
        transform=transform,
        download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=128
)

model = torchvision.models.resnet50()
model = model.to(memory_format=torch.channels_last)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = LR, momentum=0.9)
model.train()
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    with torch.autocast(device_type='cpu', dtype=torch.bfloat16):
    # with torch.cpu.amp.autocast():
        # Setting memory_format to torch.channels_last could improve performance with 4D input data. This is optional.
        data = data.to(memory_format=torch.channels_last)
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
    optimizer.step()
    print(batch_idx)
</code></pre>
<p>I tried to debug it with gdb.<br>
Error:</p>
<pre><code class="lang-auto">--Type &lt;RET&gt; for more, q to quit, c to continue without paging--

Thread 110 "python" received signal SIGSEGV, Segmentation fault.
[Switching to Thread 0x7f4d43875700 (LWP 38270)]
0x00007f4d07f26e90 in dnnl::impl::cpu::x64::brgemm_inner_product_bwd_weights_t&lt;(dnnl::impl::cpu::x64::cpu_isa_t)79&gt;::compute_diff_weights_and_bias(dnnl::impl::cpu::x64::brgemm_inner_product_bwd_weights_t&lt;(dnnl::impl::cpu::x64::cpu_isa_t)79&gt;::thread_info_t const*) const::{lambda(int, int, int)#2}::operator()(int, int, int) const ()
   from /root/anaconda3/envs/nanoPytorch1.12/lib/python3.7/site-packages/intel_extension_for_pytorch/lib/libintel-ext-pt-cpu.so
</code></pre>
<p>The <a href="https://pytorch.org/docs/stable/amp.html?highlight=autocast#torch.autocast" rel="noopener nofollow ugc">doc</a> says that  <code>torch.autocast("cpu", args...)</code> is equivalent to <code>torch.cpu.amp.autocast(args...)</code>, but <code>with torch.cpu.amp.autocast()</code> works well.<br>
So my question is that is there any difference between <code>torch.cpu.amp.autocast(..)</code> and <code>torch.autocast(..)</code></p>

ptrblck: <p>Could you test if this code still fails in the latest PyTorch release and if so, create an issue on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a>, please?</p>

y199387: <p>Thank you for your timely reply!<br>
The segmentation fault disappears and the code works well in PyTorch1.12.1+cpu.</p>