Updated training test with small model

Author: C-Achard

napari_cellseg3d/code_models/model_workers.py

@@ -1336,6 +1336,7 @@ def train(self):
 
             if model_name == "test":
                 self.quit()
+                yield TrainingReport(False)
 
             for epoch in range(self.config.max_epochs):
                 # self.log("\n")

napari_cellseg3d/code_models/models/model_test.py

@@ -13,7 +13,7 @@ def __init__(self):
         self.linear = nn.Linear(1, 1)
 
     def forward(self, x):
-        return self.linear(x)
+        return self.linear(torch.tensor(x, requires_grad=True))
 
     def get_net(self):
         return self
@@ -24,13 +24,13 @@ def get_output(self, _, input):
     def get_validation(self, val_inputs):
         return val_inputs
 
-if __name__ == "__main__":
-
-    model = TestModel()
-    model.train()
-    model.zero_grad()
-    from napari_cellseg3d.config import WEIGHTS_DIR
-    torch.save(
-        model.state_dict(),
-        WEIGHTS_DIR + f"/{get_weights_file()}"
-    )
+# if __name__ == "__main__":
+#
+#     model = TestModel()
+#     model.train()
+#     model.zero_grad()
+#     from napari_cellseg3d.config import WEIGHTS_DIR
+#     torch.save(
+#         model.state_dict(),
+#         WEIGHTS_DIR + f"/{get_weights_file()}"
+#     )