fix

configuration_ministu.py

@@ -53,4 +53,4 @@ class MiniSTUConfig(PretrainedConfig):
         self.theta = theta
         self.use_alibi = use_alibi
         self.torch_dtype = torch_dtype
-        self.device = device
+        self.device = self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # Store as string

modeling_ministu.py

@@ -434,6 +434,7 @@ class MiniSTU(PreTrainedModel):
             K=config.num_eigh,
             use_hankel_L=config.use_hankel_L,
             device=config.device,
+            dtype=config.torch_dtype,
         )
 
         self.num_layers = config.num_layers