Added new training scripts

src/trainers/quantile_trainer.py

@@ -10,12 +10,16 @@ import matplotlib.pyplot as plt
 
 
 def sample_from_dist(quantiles, output_values):
-    # both to numpy
-    quantiles = quantiles.cpu().numpy()
+    # check if tensor:
+    if isinstance(quantiles, torch.Tensor):
+        quantiles = quantiles.cpu().numpy()
 
     if isinstance(output_values, torch.Tensor):
         output_values = output_values.cpu().numpy()
 
+    if isinstance(quantiles, list):
+        quantiles = np.array(quantiles)
+
     reshaped_values = output_values.reshape(-1, len(quantiles))
 
     uniform_random_numbers = np.random.uniform(0, 1, (reshaped_values.shape[0], 1000))
@@ -60,22 +64,18 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
         data_processor: DataProcessor,
         quantiles: list,
         device: torch.device,
-        clearml_helper: ClearMLHelper = None,
         debug: bool = True,
     ):
 
         self.quantiles = quantiles
-        quantiles_tensor = torch.tensor(quantiles)
-        quantiles_tensor = quantiles_tensor.to(device)
 
-        criterion = PinballLoss(quantiles=quantiles_tensor)
+        criterion = PinballLoss(quantiles=quantiles)
         super().__init__(
             model=model,
             optimizer=optimizer,
             criterion=criterion,
             data_processor=data_processor,
             device=device,
-            clearml_helper=clearml_helper,
             debug=debug,
         )
 
@@ -252,7 +252,7 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
     def plot_quantile_percentages(
         self, task, data_loader, train: bool = True, iteration: int = None
     ):
-        quantiles = self.quantiles.cpu().numpy()
+        quantiles = self.quantiles
         total = 0
         quantile_counter = {q: 0 for q in quantiles}