Update collate_functions.py for pytorch 1.9

808e93c5 · yizhou-wang · a2321e7f · 808e93c5
Commit 808e93c5 authored Nov 09, 2021 by yizhou-wang
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

rodnet/datasets/collate_functions.py rodnet/datasets/collate_functions.py +6 -1

No files found.
--- a/rodnet/datasets/collate_functions.py
+++ b/rodnet/datasets/collate_functions.py
 import numpy as np
 import torch
 import re
-from torch._six import container_abcs, string_classes, int_classes
+try:
+    from torch._six import container_abcs, string_classes, int_classes
+except:
+    import collections.abc as container_abcs
+    string_classes = str
+    int_classes = int
 np_str_obj_array_pattern = re.compile(r'[SaUO]')
 default_collate_err_msg_format = (