import torch
import torch.nn as nn
import timm
from xgboost import XGBClassifier
model_name = 'swin_base_patch4_window7_224'

out_dim    = 1


class get_model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = timm.create_model(model_name, pretrained=False)
        self.model.head = nn.Sequential(nn.Linear(56,768),
                                              nn.Linear(768,256))
        self.last = nn.Linear(256, 128)
        self.depth1 = nn.Linear(128,64)
        self.depth2 = nn.Linear(64,6)
    def forward(self, features,targets):
        x = self.model.head(features)
        x = self.last(x)
        x = self.depth1(x)
        x = self.depth2(x)
        clf = XGBClassifier(tree_method='gpu_hist',
                             num_class = 6,objective = 'multi:softprob')
        clf.fit(x.detach().cpu().numpy(),targets.detach().cpu().numpy())
        preds = clf.predict(x.detach().cpu().numpy())
        preds = torch.tensor(preds)
        return preds

Everything is on GPU … Anything that maybe wrong was this model!!  Am i missing something?