How can i use my own GPU to train

[Dan-H-Muniz-Sanchez]

im not sure if by default u2net_train.py use my GPU or not , how can i be sure? and if the train not use the GPU cow can i change that? Thanks btw.

[ZzrMax]

you can use torch.device to check whether the train.py use GPU or not ( personal view )

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

with this function, you can check whether you use GPU. (maybe it cannot work. if not work ,you can check the instruction of torch.device )