adapt batch-size for inverse deformations

3dem · Oct 18, 2023 · 4cb7696 · 4cb7696
1 parent 34825fa
commit 4cb7696
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/dynamight/deformations/optimize_deformations.py b/dynamight/deformations/optimize_deformations.py
@@ -1314,7 +1314,8 @@ def optimize_deformations(
                                   'dec_half2_optimizer': dec_half2_optimizer.state_dict(),
                                   'indices_half1': half1_indices,
                                   'indices_val': val_indices,
-                                  'refinement_directory': refinement_star_file}
+                                  'refinement_directory': refinement_star_file,
+                                  'batch_size': batch_size}
                     if final > finalization_epochs or epoch == n_epochs-1:
                         checkpoint_file = checkpoints_directory / 'checkpoint_final.pth'
                         torch.save(checkpoint, checkpoint_file)

diff --git a/dynamight/inverse_deformations/optimize_inverse_deformations.py b/dynamight/inverse_deformations/optimize_inverse_deformations.py
@@ -56,6 +56,7 @@ def optimize_inverse_deformations(
         encoder_half2 = checkpoint['encoder_half2']
         decoder_half1 = checkpoint['decoder_half1']
         decoder_half2 = checkpoint['decoder_half2']
+        batch_size = checkpoint['batch_size']
         poses = checkpoint['poses']
 
         encoder_half1.load_state_dict(checkpoint['encoder_half1_state_dict'])

diff --git a/dynamight/models/__pycache__/decoder.cpython-310.pyc b/dynamight/models/__pycache__/decoder.cpython-310.pyc