@ptrblck thank you for the quick reply, I didn’t push data into GPU in Dataset.
This is my Dataset:
import os
import cv2
import numpy as np 
import pandas as pd

import torch
import torch.nn as nn
from torch.utils.data import Dataset

from matplotlib import pyplot as plt
import albumentations.pytorch as AT
from albumentations import (
    RandomRotate90, Flip, Transpose, GaussNoise, Blur, VerticalFlip, HorizontalFlip,  \
    HueSaturationValue, RGBShift, RandomBrightness, Resize, Normalize, Compose, CenterCrop)
from PIL import Image

def transforms(size_image):
    return Compose([
        Resize(height = size_image[0], width = size_image[1]),
        # Red - Green - Blue right now
        Normalize(mean=(0.406, 0.515, 0.323), std=(0.195, 0.181, 0.178)),
        AT.ToTensor()
    ])

def augmentation(size_image,p=0.5):
    return Compose([
        RandomRotate90(),
        Flip(),
        Transpose(),
        GaussNoise(),
        Blur(),
        VerticalFlip(),
        HorizontalFlip(),
        HueSaturationValue(hue_shift_limit=5, sat_shift_limit=15, val_shift_limit=10),
        RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10),
        RandomBrightness(limit = 0.05),
        CenterCrop(height = 150, width = 150, p = 0.5)    
    ], p=p)

class PyTorchImageDataset(Dataset):
    def __init__(self, image_list,train, labels, size_image, **kwags):
        self.image_list = image_list
        self.transforms = transforms(size_image)
        self.labels = labels
        self.train = train
        self.augment = augmentation(size_image)
        import json
        with open('C:/Users/user/name2path.json', 'r') as f: 
            self.image_path_dir = json.load(f)

    def __len__(self):
        return len(self.image_list)
    
    def __getitem__(self, i):
        image_path = self.image_path_dir[self.image_list[i]]
        image = np.array(Image.open(image_path))
        label = self.labels[i]
        if self.train:
            image = self.augment(image=image)['image']   
        image = self.transforms(image = image)['image']
        return image, label

    def isImage(self,path):
        all_image_ext =  ["jpg","gif","png","tga","jpeg"]
        return True if (path.split('.')[-1].lower()) in all_image_ext else False

And my train code:
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
import sys

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader

from models import se_resnext50_32x4d
from tqdm import tqdm

from sklearn.model_selection import train_test_split
from utilities import *

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
params = {'workers': 4,
            'batch_size': 16,
            'num_epochs': 100,
            'lr': 0.001,
            'size_image': [480, 768],
            'checkpoint': True,
            'save_path': 'C:/Users/user/save',
            'dataFrame_path': 'C:/Users/user/train.csv',
            'max_patience': 5}

checkDir(params['save_path'])
df = pd.read_csv(params['dataFrame_path'])


IDs = df.id.to_list()
labels = df.landmark_id.to_list()
train ,val, y_train,y_val = train_test_split(IDs, labels, test_size = 0.2, random_state = 42, shuffle = True)

train_dataset = PyTorchImageDataset(image_list = train, 
                                    labels = y_train, \
                                    train = True, 
                                    size_image = params['size_image'])
train_dataloader = DataLoader(dataset = train_dataset, 
                                batch_size = params['batch_size'], \
                                shuffle = True,
                                num_workers=params['workers'])

val_dataset = PyTorchImageDataset(image_list=val,
                                    labels=y_val, \
                                    train = False, \
                                    size_image = params['size_image'])
val_dataloader = DataLoader(dataset=val_dataset, \
                            batch_size=params['batch_size'], \
                            shuffle=True, \
                            num_workers=params['workers'])

dataloader = {'train': train_dataloader, 'valid': val_dataloader}

model = se_resnext50_32x4d()
model.to(device)
model.train()

criterion = nn.CrossEntropyLoss()
softmax = nn.Softmax()
optimizer = optim.Adam(model.parameters(), lr=params['lr'], betas=(0.9, 0.999))

patience = 0
last_acc = -1
best_acc = 0.3

for epoch in range(params['num_epochs']):
    for phase in ['train','valid']:
        if phase == 'train':
            running_loss = 0
            for i, data in enumerate(tqdm(dataloader[phase])):
                optimizer.zero_grad()
                data_batch = data[0].to(device)
                b_size = data_batch.size(0)                
                label = data[1].type(torch.long)
                label = label.to(device)
                output = model(data_batch)
                prob = softmax(output)
                loss = criterion(output, label)
                running_loss+= loss.item()
                loss.backward()
                optimizer.step()  
            print('epoch %d train loss: %.3f' %(epoch + 1, float(running_loss)/(1+i)))

        else:
            running_loss = 0
            with torch.no_grad():
                for i, data in enumerate(tqdm(dataloader[phase])):
                    data_batch = data[0].to(device)
                    b_size = data_batch.size(0)
                    label = data[1].type(torch.long)
                    label = label.to(device)
                    output = model(data_batch)
                    prob = softmax(output)
                    loss = criterion(output, label)
                
                    if phase == 'valid' and i == 0:

                        valid_label = label.cpu().detach().numpy()
                        
                        valid_prob = prob.cpu().detach().numpy()
                                                    
                        running_loss+= loss.item()
                    else:
                        torch.cuda.synchronize() 

                        temp_label = label.cpu().detach().numpy()

                        valid_label = np.concatenate((valid_label, temp_label))

                        torch.cuda.synchronize() 

                        temp_prob = prob.cpu().detach().numpy()

                        valid_prob = np.concatenate((valid_prob, temp_prob), axis = 0)

                        running_loss+= loss.item()
    if phase == 'valid':
        
        last = 0
        torch.cuda.synchronize() 
        predict_labels = np.argmax(valid_prob, axis=1)
        acc = accuracy(valid_label, predict_labels)
        loss = round(float(running_loss)/(i + 1), 4)
        print('epoch %d valid acc: %.3f' %(epoch + 1, acc))
        print('epoch %d valid loss: %.3f' %(epoch + 1, loss))
        
        path_name = params['save_path'] + 'se_resnext_' + str(last + epoch +  1) + '_loss_' + str(loss) + '_acc_' + str(acc) + '.pth'
        if acc > best_acc:
            torch.save(model.state_dict(), path_name)
            best_acc = acc
        if acc < last_acc:
            patience+=1
        elif acc >= last_acc:
            patience = 0
        if patience == params['max_patience']:
            torch.save(model.state_dict(), path_name)
            sys.exit()
        last_acc = acc