diff --git a/networks/med.py b/networks/med.py
index 7b00a35..ce51297 100644
--- a/networks/med.py
+++ b/networks/med.py
@@ -36,12 +36,14 @@
     SequenceClassifierOutput,
     TokenClassifierOutput,
 )
-from transformers.modeling_utils import (
-    PreTrainedModel,
+from transformers.modeling_utils import PreTrainedModel
+
+from transformers.pytorch_utils import (
     apply_chunking_to_forward,
     find_pruneable_heads_and_indices,
     prune_linear_layer,
 )
+
 from transformers.utils import logging
 from transformers.models.bert.configuration_bert import BertConfig
 
diff --git a/networks/nlvr_encoder.py b/networks/nlvr_encoder.py
index 1946bb4..55c88b6 100644
--- a/networks/nlvr_encoder.py
+++ b/networks/nlvr_encoder.py
@@ -26,12 +26,14 @@
     SequenceClassifierOutput,
     TokenClassifierOutput,
 )
-from transformers.modeling_utils import (
-    PreTrainedModel,
+from transformers.modeling_utils import PreTrainedModel
+
+from transformers.pytorch_utils import (
     apply_chunking_to_forward,
     find_pruneable_heads_and_indices,
     prune_linear_layer,
 )
+
 from transformers.utils import logging
 from transformers.models.bert.configuration_bert import BertConfig