config.py

import torch

class Config:
    def __init__(self):
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        self.is_half = self.device != "cpu"
        self.n_cpu = 0
        self.gpu_name = None
        self.gpu_mem = None
        self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()

    def device_config(self) -> tuple:
        if torch.cuda.is_available():
            i_device = int(self.device.split(":")[-1])
            self.gpu_name = torch.cuda.get_device_name(i_device)
            if ("16" in self.gpu_name and "V100" not in self.gpu_name):
                print("16-series GPU detected, forcing full precision (half-precision not supported).")
                self.is_half = False
            self.gpu_mem = int(torch.cuda.get_device_properties(i_device).total_memory / 1024 / 1024 / 1024 + 0.4)
            if self.gpu_mem <= 4:
                x_pad, x_query, x_center, x_max = 1, 5, 5, 8
            elif self.gpu_mem <= 5:
                x_pad, x_query, x_center, x_max = 1, 6, 6, 8
            else:
                x_pad, x_query, x_center, x_max = 3, 8, 8, 12
        else:
            x_pad, x_query, x_center, x_max = 1, 5, 5, 8
        return x_pad, x_query, x_center, x_max