Refactor data module

[ilyes319]

No description provided.

[LarsSchaaf]

Currently doesn't support wandb, as gives this error:

TypeError: Object of type KeySpecification is not JSON serializable

From line:

args_dict_json = json.dumps(args_dict)

One solution:

class CustomEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, KeySpecification): 
            return obj.__dict__  
        return super().default(obj)

args_dict_json = json.dumps(args_dict, cls=CustomEncoder)