Argumente Subklassen

models.py

@@ -107,9 +107,9 @@ class SummarisationModel(nn.Module):
 
 class ActorOnlySummarisationModel(SummarisationModel):
     
-    def __init__(self):
+    def __init__(self, gpu):
 
-        super().__init__()
+        super().__init__(gpu)
 
         self.optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
 
@@ -162,9 +162,9 @@ class ActorOnlySummarisationModel(SummarisationModel):
 
 class SummarisationModelWithCrossEntropyLoss(SummarisationModel):
     
-    def __init__(self):
+    def __init__(self, gpu):
 
-        super().__init__()
+        super().__init__(gpu)
 
         self.loss_fn = nn.BCELoss(reduction='sum')
         self.optimizer = torch.optim.Adam(self.parameters(), lr=0.001)