Collect all contrib models in __init__ and add unit tests for init

qlib/contrib/model/pytorch_gru.py

@@ -8,6 +8,7 @@ from __future__ import print_function
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
 from ...utils import (
     unpack_archive_with_buffer,
@@ -273,11 +274,11 @@ class GRU(Model):
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         index = x_test.index
         self.gru_model.eval()
         x_values = x_test.values