kernel-elastic-autoencoder 3.3.0__tar.gz → 3.3.1__tar.gz

{kernel_elastic_autoencoder-3.3.0 → kernel_elastic_autoencoder-3.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.3.0
+Version: 3.3.1
 Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
 License: MIT
 Author: Felix Rotter-McCartney

{kernel_elastic_autoencoder-3.3.0 → kernel_elastic_autoencoder-3.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.3.0"
+version = "3.3.1"
 description = "Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)"
 authors = [
     { name = "Felix Rotter-McCartney", email = "felix.rotter@mail.utoronto.ca" }

{kernel_elastic_autoencoder-3.3.0 → kernel_elastic_autoencoder-3.3.1}/src/kernel_elastic_autoencoder/training.py

@@ -123,7 +123,7 @@ class Trainer:
         accelerator.wait_for_everyone()
         for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
             model.train()
-            train_loss = torch.tensor([])
+            train_loss = torch.tensor([], device=accelerator.device)
             for input_ids, conditions, token_mask, condition_mask in tqdm(
                 dataloader_train, desc=f"Epoch {epoch}, Train Batch"
             ):
@@ -140,7 +140,7 @@ class Trainer:
             accelerator.print(f"Avg. train loss: {train_loss.mean().item()}")
 
             model.eval()
-            test_loss = torch.tensor([])
+            test_loss = torch.tensor([], device=accelerator.device)
             for input_ids, conditions, token_mask, condition_mask in tqdm(
                 dataloader_test, desc=f"Epoch {epoch}, Test Batch"
             ):