I have this problem and I don’t understand why. If I try to use EfficientNetB3, I get this error that is due to the line 92 as you can see below.
But the thing is… if I try to use EfficientNetB0 or B1 for examples, I simply change the Imagesize and it works!!! Why am I getting this error so??? What is the problem through EfficientNetB3 and my code?

Can be the Image Size the problem? Because all EfficientNet B0, B1, B2 have Imagesize below 300, while B3 and so on, has >= 300, and in fact from B3 onward the system doesn’t work

The error:

Downloading: "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pth" to /user/sfasulo/.cache/torch/hub/checkpoints/efficientnet-b3-5fb5a3c3.pth
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 47.1M/47.1M [00:01<00:00, 42.5MB/s]
Loaded pretrained weights for efficientnet-b3
Epoch 1/20
Traceback (most recent call last):
  File "/user/sfasulo/p.py", line 129, in <module>
    train(model, train_loader, criterion, optimizer, DEVICE)
  File "/user/sfasulo/p.py", line 92, in train
    outputs = model(images)
              ^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/efficientnet_pytorch/model.py", line 314, in forward
    x = self.extract_features(inputs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/efficientnet_pytorch/model.py", line 296, in extract_features
    x = block(x, drop_connect_rate=drop_connect_rate)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/efficientnet_pytorch/model.py", line 111, in forward
    x = self._swish(x)
        ^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/efficientnet_pytorch/utils.py", line 80, in forward
    return SwishImplementation.apply(x)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/torch/autograd/function.py", line 598, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/sfasulo/miniconda3/lib/python3.12/site-packages/efficientnet_pytorch/utils.py", line 67, in forward
    result = i * torch.sigmoid(i)
                 ^^^^^^^^^^^^^^^^
RuntimeError: NVML_SUCCESS == r INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1716905969118/work/c10/cuda/CUDACachingAllocator.cpp":844, please report a bug to PyTorch. 

My code is the following:

ROOT_DIR_CLASS = "/path/classifier_no300/"
IMG_SIZE = 300
BATCH_SIZE = 64
NUM_CLASSES = 3
EPOCHS = 20
LEARNING_RATE = 1e-3
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Custom Dataset
class CustomDataset(Dataset):
    def __init__(self, dataframe, transform=None):
        self.dataframe = dataframe
        self.transform = transform

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        img_name = os.path.join(ROOT_DIR_CLASS, self.dataframe.iloc[idx, 1])
        image = cv2.imread(img_name)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        label = self.dataframe.iloc[idx, 0]
        if self.transform:
            image = self.transform(image)
        return image, label

# Data Augmentation and Normalization
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Build a dataframe
class_labels = []
for item in os.listdir(ROOT_DIR_CLASS):
    all_classes = os.listdir(ROOT_DIR_CLASS + '/' + item)
    for i in all_classes:
        class_labels.append((item, str(item) + '/' + i))

df = pd.DataFrame(data=class_labels, columns=['Labels', 'image'])

# Label encoding
label_encoder = LabelEncoder()
df['Labels'] = label_encoder.fit_transform(df['Labels'])

# Split the dataset
train_df, test_df = train_test_split(df, test_size=0.05, random_state=415)

# Create datasets
train_dataset = CustomDataset(train_df, transform=transform)
test_dataset = CustomDataset(test_df, transform=transform)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# Load the pre-trained EfficientNetB3 model
model = EfficientNet.from_pretrained('efficientnet-b3')
model._fc = nn.Linear(model._fc.in_features, NUM_CLASSES)
model = model.to(DEVICE)

# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Training loop
def train(model, train_loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    print(f"Training Loss: {epoch_loss:.4f}")

# Evaluation loop
def evaluate(model, test_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item()

            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / len(test_loader)
    accuracy = 100 * correct / total
    print(f"Validation Loss: {epoch_loss:.4f}, Accuracy: {accuracy:.2f}%")

# Main training and evaluation loop
for epoch in range(EPOCHS):
    print(f"Epoch {epoch+1}/{EPOCHS}")
    train(model, train_loader, criterion, optimizer, DEVICE)
    evaluate(model, test_loader, criterion, DEVICE)

# Save the model
#torch.save(model.state_dict(), '/path/efficientnet_b3_finetuned_1.pth')