Synthetic Dataset - ILC Performing Worse

[amimem]

We have to find out why ilc is performing worse on the synthetic dataset:

• [x] Try to recreate the graph in GB's notebook using the JAX code
  • [x] If the dataset is the problem, change it and find out the statistical variables that cause this
  • [x] Otherwise, fix the code

[amimem]

Update:

• JAX is working fine

New Task:

• [ ] Changing JAX code to work with a list of datasets or dataloaders