optimized the imports

test/modules/test_distribution_output.py

@@ -1,13 +1,9 @@
-import pytest
-from typing import Iterable, List, Tuple
+from typing import List, Tuple
 
 import numpy as np
-
+import pytest
 import torch
 import torch.nn as nn
-from torch.nn.utils import clip_grad_norm_
-from torch.utils.data import TensorDataset, DataLoader
-from torch.optim import SGD
 from torch.distributions import (
     StudentT,
     Beta,
@@ -17,6 +13,9 @@ from torch.distributions import (
     Independent,
     Normal,
 )
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import SGD
+from torch.utils.data import TensorDataset, DataLoader
 
 from pts.modules import (
     DistributionOutput,

test/modules/test_feature.py

@@ -1,10 +1,8 @@
-import pytest
 from itertools import chain, combinations
 
-
+import pytest
 import torch
 import torch.nn as nn
-from torch.distributions import Uniform
 
 from pts.modules import FeatureEmbedder, FeatureAssembler
 

test/modules/test_scaler.py

@@ -11,14 +11,12 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-import pytest
 import numpy as np
-
+import pytest
 import torch
 
 from pts.modules import MeanScaler, NOPScaler
 
-
 test_cases = [
     (
         MeanScaler(),