Update TabPFN/layer.py

TabPFN/layer.py

@@ -1,8 +1,16 @@
 from functools import partial
 
 from torch import nn
-from torch.nn.modules.transformer import *
-from torch.nn.modules.transformer import _get_activation_fn
+from torch.nn.modules.transformer import (
+    _get_activation_fn,
+    Module,
+    Tensor,
+    Optional,
+    MultiheadAttention,
+    Linear,
+    Dropout,
+    LayerNorm,
+)
 
 from torch.utils.checkpoint import checkpoint