Sequential.patch 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. --- /usr/local/lib/python3.5/dist-packages/torch/nn/modules/container.py
  2. +++ /usr/local/lib/python3.5/dist-packages/torch/nn/modules/container.py
  3. @@ -22,14 +22,6 @@
  4. ]))
  5. """
  6. - @overload
  7. - def __init__(self, *args: Module) -> None:
  8. - ...
  9. -
  10. - @overload
  11. - def __init__(self, arg: 'OrderedDict[str, Module]') -> None:
  12. - ...
  13. -
  14. def __init__(self, *args):
  15. super(Sequential, self).__init__()
  16. if len(args) == 1 and isinstance(args[0], OrderedDict):
  17. @@ -39,7 +31,7 @@
  18. for idx, module in enumerate(args):
  19. self.add_module(str(idx), module)
  20. - def _get_item_by_idx(self, iterator, idx) -> T:
  21. + def _get_item_by_idx(self, iterator, idx):
  22. """Get the idx-th item of the iterator"""
  23. size = len(self)
  24. idx = operator.index(idx)
  25. @@ -48,18 +40,17 @@
  26. idx %= size
  27. return next(islice(iterator, idx, None))
  28. - @_copy_to_script_wrapper
  29. - def __getitem__(self, idx) -> Union['Sequential', T]:
  30. + def __getitem__(self, idx):
  31. if isinstance(idx, slice):
  32. return self.__class__(OrderedDict(list(self._modules.items())[idx]))
  33. else:
  34. return self._get_item_by_idx(self._modules.values(), idx)
  35. - def __setitem__(self, idx: int, module: Module) -> None:
  36. - key: str = self._get_item_by_idx(self._modules.keys(), idx)
  37. + def __setitem__(self, idx, module):
  38. + key = self._get_item_by_idx(self._modules.keys(), idx)
  39. return setattr(self, key, module)
  40. - def __delitem__(self, idx: Union[slice, int]) -> None:
  41. + def __delitem__(self, idx):
  42. if isinstance(idx, slice):
  43. for key in list(self._modules.keys())[idx]:
  44. delattr(self, key)
  45. @@ -67,26 +58,16 @@
  46. key = self._get_item_by_idx(self._modules.keys(), idx)
  47. delattr(self, key)
  48. - @_copy_to_script_wrapper
  49. - def __len__(self) -> int:
  50. + def __len__(self):
  51. return len(self._modules)
  52. - @_copy_to_script_wrapper
  53. def __dir__(self):
  54. keys = super(Sequential, self).__dir__()
  55. keys = [key for key in keys if not key.isdigit()]
  56. return keys
  57. - @_copy_to_script_wrapper
  58. - def __iter__(self) -> Iterator[Module]:
  59. - return iter(self._modules.values())
  60. -
  61. - # NB: We can't really type check this function as the type of input
  62. - # may change dynamically (as is tested in
  63. - # TestScript.test_sequential_intermediary_types). Cannot annotate
  64. - # with Any as TorchScript expects a more precise type
  65. def forward(self, input):
  66. - for module in self:
  67. + for module in self._modules.values():
  68. input = module(input)
  69. return input