gninatorch.utils module

gninatorch.utils.log_print(metrics, title: str | None = None, epoch: int | None = None, epoch_time: float | None = None, elapsed_time: float | None = None, stream=<_io.TextIOWrapper name='<stdout>' mode='w' encoding='utf-8'>)[source]

Print metrics to the console.

Parameters:

  • metrics – Dictionary of metrics

  • title (str) – Title to print

  • epoch (int) – Epoch number

  • stream – Outoput stream

gninatorch.utils.print_args(args: ~argparse.Namespace, header: str | None = None, stream=<_io.TextIOWrapper name='<stdout>' mode='w' encoding='utf-8'>)[source]

Print command line arguments to stream.py

Parameters:

  • args (argparse.Namespace) – Command line arguments

  • header (str)

  • stream – Output stream

gninatorch.utils.set_device(device_name: str) → device[source]

Set the device to use.

Parameters:

device_name (str) – Name of the device to use ("cpu", "cuda", "cuda:0", …)

Returns:

PyTorch device

Return type:

torch.device

Notes

This function also set the global device for molgrid so that the molgrid.ExampoleProvider works on the correct device.

https://github.com/gnina/libmolgrid/issues/43