Fixed two typing errors

OmerRonen · Jun 7, 2019 · 462e986 · 462e986
1 parent 91b5704
commit 462e986
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/bartpy/model.py b/bartpy/model.py
@@ -1,5 +1,5 @@
 from copy import deepcopy, copy
-from typing import List, Generator
+from typing import List, Generator, Optional
 
 import numpy as np
 import pandas as pd
@@ -13,7 +13,7 @@
 class Model:
 
     def __init__(self,
-                 data: Data,
+                 data: Optional[Data],
                  sigma: Sigma,
                  trees=None,
                  n_trees: int = 50,

diff --git a/bartpy/samplers/schedule.py b/bartpy/samplers/schedule.py
@@ -1,4 +1,4 @@
-from typing import Callable, Generator
+from typing import Callable, Generator, Text, Tuple
 
 from bartpy.model import Model
 from bartpy.samplers.leafnode import LeafNodeSampler
@@ -30,7 +30,7 @@ def __init__(self,
         self.sigma_sampler = sigma_sampler
         self.tree_sampler = tree_sampler
 
-    def steps(self, model: Model) -> Generator[Callable[[Model], Sampler], None, None]:
+    def steps(self, model: Model) -> Generator[Tuple[Text, Callable[[], float]], None, None]:
         """
         Create a generator of the steps that need to be called to complete a full Gibbs sample