Repository: pytorch
Issue: s.isIntegral(false) INTERNAL ASSERT FAILED
Labels: needs reproduction,module: autograd,triaged

### 🐛 Describe the bug

When training ViT_b_16 (https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html#torchvision.models.vit_b_16) on CUDA 

```
    model = helper.train_model(model, dataloaders, criterion, optimizer, scheduler,
  File "/home/birdy/meng_thesis/code/master_ifcb_classifier/utils/resnet_experiment_helpers.py", line 298, in train_model
    loss.backward()
  File "/home/birdy/.conda/envs/cv-env/lib/python3.10/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/birdy/.conda/envs/cv-env/lib/python3.10/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: s.isIntegral(false) INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch-base_1664259963002/work/aten/src/ATen/ScalarOps.h":39, please report a bug to PyTorch.
```

Code defining model 
```
model = models.vit_b_16(weights='IMAGENET1K_V1')
model.heads.head = nn.Linear(model.heads.head.in_features, NUM_CLASSES) 
model = nn.DataParallel(model)
model= model.to(device)
```

Code for training function 
```
def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs=25, save_checkpoints = False, DEST='', model_name="ResNet", val_fqn=1, device="cpu"):
    since = time.time()

    # Create a temporary directory to save training checkpoints
    with TemporaryDirectory() as tempdir:
        #best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')
        best_model_params_path = f'{DEST}/{model_name}-best.pt'
        torch.save(model.state_dict(), best_model_params_path)
        best_acc = 0.0
        # intialize training history if model doesn't already have it
        if hasattr(model, 'history') == False: 
            model.history = {
                'train_loss': [],
                'train_acc': [],
                'train_class_acc': [],
                'train_epoch_duration': [],
                'val_epochs': [],
                'val_loss': [],
                'val_acc': [],
                'val_class_acc': [],
                'val_epoch_duration': []
            }
        dt_string = datetime.now().strftime("%d-%m-%Y-%H-%M-%S")

        for epoch in range(num_epochs):
            print(f'Epoch {epoch}/{num_epochs - 1}')
            print('-' * 10)
            epoch_start_time = time.time()

            # Each epoch has a training and validation phase
            for phase in ['train', 'val']:
                if phase == 'val' and epoch % val_fqn !=0:
                    continue # only validate model every nth time
                if phase == 'train':
                    model.train()  # Set model to training mode
                else:
                    model.history['val_epochs'].append(epoch)
                    model.eval()   # Set model to evaluate mode

                running_loss = 0.0
                running_corrects = 0
                running_class_acc = 0.0
                batch_counter = 0
                # Iterate over data.
                for inputs, labels in tqdm(dataloaders[phase]):
                    batch_counter += 1
                    inputs = inputs.to(device)
                    labels = labels.to(device)

                    # labels = torch.tensor(labels).to(device)

                    # zero the parameter gradients
                    optimizer.zero_grad()

                    # print(inputs.shape)
                    # forward
                    # track history if only in train
                    with torch.set_grad_enabled(phase == 'train'):
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)
                        loss = criterion(outputs, labels)

                        # backward + optimize only if in training phase
                        if phase == 'train':
                            loss.backward()
                            optimizer.step()

                    # statistics
                    running_loss += loss.item() * inputs.size(0)
                    running_corrects += torch.sum(preds == labels.data)
                    avg_class_acc, _ = calculate_per_class_accuracy(labels.data.cpu().numpy(), preds.cpu().numpy())
                    running_class_acc += avg_class_acc
                if phase == 'train':
                    if type(scheduler) == lr_scheduler.ReduceLROnPlateau:
                        scheduler.step(running_loss)
                    else:
                      scheduler.step() 

                epoch_loss = running_loss / len(dataloaders[phase].dataset)
                epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
                epoch_class_acc = running_class_acc / batch_counter
                epoch_duration = time.time() - epoch_start_time

                print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
                model.history[f'{phase}_loss'].append(epoch_loss)
                model.history[f'{phase}_acc'].append(epoch_acc.item())
                model.history[f'{phase}_class_acc'].append(epoch_class_acc)
                model.history[f'{phase}_epoch_duration'].append(epoch_duration)


                if save_checkpoints:
                  # write over the old checkpoint - saves mem space
                  PATH = f"{DEST}/{model_name}-{dt_string}-checkpoint.pt"
                  save_checkpoint(model, optimizer, PATH, epoch)


                # deep copy the model
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    torch.save(model.state_dict(), best_model_params_path)

            print()

        time_elapsed = time.time() - since
        model.history['time_elapsed'] = time_elapsed
        print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
        print(f'Best val Acc: {best_acc:4f}')

        # load best model weights
        # model.load_state_dict(torch.load(best_model_params_path))
    return model
```

### Versions


Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.3 (Ootpa) (ppc64le)
GCC version: (GCC) 8.3.1 20191121 (Red Hat 8.3.1-5)
Clang version: Could not collect
CMake version: version 3.11.4
Libc version: glibc-2.28

Python version: 3.10.13 (main, Sep 11 2023, 13:14:29) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.18.0-240.el8.ppc64le-ppc64le-with-glibc2.28
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

CPU:
Architecture:        ppc64le
Byte Order:          Little Endian
CPU(s):              128
On-line CPU(s) list: 0-127
Thread(s) per core:  4
Core(s) per socket:  16
Socket(s):           2
NUMA node(s):        2
Model:               2.2 (pvr 004e 1202)
Model name:          POWER9, altivec supported
CPU max MHz:         3800.0000
CPU min MHz:         2166.0000
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            10240K
NUMA node0 CPU(s):   0-63
NUMA node8 CPU(s):   64-127

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.12.1
[pip3] torchprofile==0.0.4
[pip3] torchvision==0.13.1
[conda] _pytorch_select           2.0                      cuda_2    https://opence.mit.edu
[conda] cudatoolkit               11.4.4              h5d40d8d_10    https://opence.mit.edu
[conda] cudnn                     8.3.0_11.4           h5b181e6_1    https://opence.mit.edu
[conda] nccl                      2.12.7               cuda11.4_1    https://opence.mit.edu
[conda] numpy                     1.23.5          py310h87cc683_0
[conda] numpy-base                1.23.5          py310hac71eb6_0
[conda] pytorch-base              1.12.1          cuda11.4_py310_pb3.19_2    https://opence.mit.edu
[conda] torchprofile              0.0.4                    pypi_0    pypi
[conda] torchvision               0.13.1          cuda11.4_py310_1    https://opence.mit.edu
[conda] torchvision-base          0.13.1          cuda11.4_py310_1    https://opence.mit.edu


cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan