model = SuperPointNet_gauss2()
    
for param in model.parameters():
    param.requires_grad = False
for param in model.linear.parameters(): # 指定訓練
    param.requires_grad = True

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

# check keras-like model summary using torchsummary
from torchsummary import summary
summary(model, input_size=(1, 240, 320))