From 1147d2e4dee43c647b4baf1eeddbced1e2dd5010 Mon Sep 17 00:00:00 2001
From: Bruno Rousseau <rousseau.bruno@gmail.com>
Date: Fri, 5 Apr 2024 16:01:20 -0400
Subject: [PATCH] Updated example.

---
 examples/local/config_diffusion.yaml | 11 +++++++----
 examples/local/run_diffusion.sh      |  2 --
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/examples/local/config_diffusion.yaml b/examples/local/config_diffusion.yaml
index 8ed5fa27..0ba99c24 100644
--- a/examples/local/config_diffusion.yaml
+++ b/examples/local/config_diffusion.yaml
@@ -1,6 +1,6 @@
 # general
-loss: cross_entropy
-max_epoch: 5
+max_epoch: 50
+log_every_n_steps: 1
 exp_name: exp_example
 # set to null to avoid setting a seed (can speed up GPU computation, but
 # results will not be reproducible)
@@ -8,7 +8,7 @@ seed: 1234
 
 # data
 data:
-  batch_size: 32
+  batch_size: 1024
   num_workers: 0
   max_atom: 512
 
@@ -31,4 +31,7 @@ optimizer:
 early_stopping:
   metric: val_loss
   mode: min
-  patience: 3
\ No newline at end of file
+  patience: 3
+
+logging:
+  - tensorboard
diff --git a/examples/local/run_diffusion.sh b/examples/local/run_diffusion.sh
index 8e1d3a7e..2a1ac9a2 100755
--- a/examples/local/run_diffusion.sh
+++ b/examples/local/run_diffusion.sh
@@ -1,6 +1,5 @@
 #!/bin/bash
 
-LOG=debug.log
 CONFIG=config_diffusion.yaml
 DATA_DIR=../../data/si_diffusion_v1
 PROCESSED_DATA=${DATA_DIR}/processed
@@ -8,7 +7,6 @@ DATA_WORK_DIR=./tmp_work_dir/
 OUTPUT=debug
 
 python ../../crystal_diffusion/train_diffusion.py \
-    --log $LOG \
     --config $CONFIG \
     --data $DATA_DIR \
     --processed_datadir $PROCESSED_DATA \