Bettr environment, lower forward samples.

usnistgov · Jun 23, 2024 · e187062 · e187062
1 parent 13da6ed
commit e187062
Show file tree

Hide file tree

Showing 3 changed files with 148 additions and 317 deletions.
diff --git a/atomgpt/examples/forward_model/config.json b/atomgpt/examples/forward_model/config.json
@@ -2,18 +2,18 @@
     "id_prop_path": "atomgpt/examples/forward_model/id_prop.csv",
     "prefix": "atomgpt_run",
     "model_name": "gpt2",
-    "batch_size": 5,
+    "batch_size": 2,
     "max_length": 512,
-    "num_epochs": 3,
+    "num_epochs": 1,
     "latent_dim": 1024,
     "learning_rate": 0.001,
     "test_each_run": true,
     "include_struct": false,
     "pretrained_path": "",
     "seed_val": 42,
-    "n_train": 40,
-    "n_val": 5,
-    "n_test": 5,
+    "n_train": 2,
+    "n_val": 2,
+    "n_test": 2,
     "output_dir": "atomgpt/examples/forward_model/out",
     "train_ratio": null,
     "val_ratio": 0.1,

diff --git a/atomgpt/inverse_models/kernels/__init__.py b/atomgpt/inverse_models/kernels/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .cross_entropy_loss import fast_cross_entropy_loss
+from atomgpt.inverse_models.kernels.cross_entropy_loss import fast_cross_entropy_loss
 from .rms_layernorm import fast_rms_layernorm
 from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
 from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel