Minor fix to type annotations.

PiperOrigin-RevId: 328491975
victor-mageto · Aug 26, 2020 · 9bd09fc · 9bd09fc
1 parent 486d8ca
commit 9bd09fc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rl_unplugged/dm_control_suite.py b/rl_unplugged/dm_control_suite.py
@@ -30,7 +30,7 @@
 import collections
 import functools
 import os
-from typing import Dict, Tuple, Set
+from typing import Dict, Optional, Tuple, Set
 
 from acme import wrappers
 from dm_control import composer
@@ -778,7 +778,7 @@ def dataset(root_path: str,
             shapes: Dict[str, Tuple[int]],
             num_threads: int,
             batch_size: int,
-            uint8_features: Set[str] = None,
+            uint8_features: Optional[Set[str]] = None,
             num_shards: int = 100,
             shuffle_buffer_size: int = 100000,
             sarsa: bool = True) -> tf.data.Dataset: