run pre-commit

Author: sotagg

pytorch/skorch_simple.py

@@ -15,14 +15,16 @@
 
 import numpy as np
 import optuna
+from optuna.integration import SkorchPruningCallback
 import skorch
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from optuna.integration import SkorchPruningCallback
+from torchvision import datasets
+
 from sklearn.metrics import accuracy_score
 from sklearn.model_selection import train_test_split
-from torchvision import datasets
+
 
 SUBSET_RATIO = 0.4
 
@@ -39,9 +41,7 @@
 X = X[indices][:N]
 y = y[indices][:N]
 
-X_train, X_test, y_train, y_test = train_test_split(
-    X, y, test_size=0.25, random_state=42
-)
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -96,9 +96,7 @@ def objective(trial: optuna.Trial) -> float:
     )
     args = parser.parse_args()
 
-    pruner = (
-        optuna.pruners.MedianPruner() if args.pruning else optuna.pruners.NopPruner()
-    )
+    pruner = optuna.pruners.MedianPruner() if args.pruning else optuna.pruners.NopPruner()
 
     study = optuna.create_study(direction="maximize", pruner=pruner)
     study.optimize(objective, n_trials=100, timeout=600)