Repository: mit-ll-responsible-ai/hydra-zen (python)

Issue:
Add warning to PL example about issues of keeping models in-memory
Regarding: [Run Boilerplate-Free ML Experiments with PyTorch Lightning & hydra-zen](https://mit-ll-responsible-ai.github.io/hydra-zen/how_to/pytorch_lightning.html)

While the following

```python

    # return the trained model instance and the final fit
    return (
        lit_module,
        final_fit.detach().numpy().ravel(),
    )
```

is very convenient and appropriate for our basic how-to, it is inadvisable to return the model in-memory at the end of the job. A multirun will result in multiple models being held in-memory, which could easily lead to memory issues at-scale.

We should add a warning.

Pull request:
PL How-To revision
Closes #152 

Patch:
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 5f23a1a1..0005f5a2 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -41,7 +41,7 @@ jobs:
         restore-keys: |
           ${{ runner.os }}-node-
     - name: Install pyright
-      run: sudo npm install -g pyright@">=1.1.183"
+      run: sudo npm install -g pyright@">=1.1.148"
     - uses: actions/checkout@v1
     - name: Set up Python
       uses: actions/setup-python@v2
diff --git a/.github/workflows/tox_run.yml b/.github/workflows/tox_run.yml
index 634df065..f7f47427 100644
--- a/.github/workflows/tox_run.yml
+++ b/.github/workflows/tox_run.yml
@@ -159,7 +159,7 @@ jobs:
         restore-keys: |
           ${{ runner.os }}-node-
     - name: Install pyright
-      run: sudo npm install -g pyright@">=1.1.183"
+      run: sudo npm install -g pyright@">=1.1.148"
     - uses: actions/checkout@v1
     - name: Set up Python
       uses: actions/setup-python@v2
diff --git a/docs/source/how_to/pytorch_lightning.rst b/docs/source/how_to/pytorch_lightning.rst
index e11587bf..0330d1ff 100644
--- a/docs/source/how_to/pytorch_lightning.rst
+++ b/docs/source/how_to/pytorch_lightning.rst
@@ -107,9 +107,9 @@ following code. Here, we define our single-layer neural network and the `lightni
 Creating Our Configs and Task Function
 ======================================
 
-Create another script - named ``experiment.py`` - in the same directory as ``zen_model.py``. Here, we will create the configs for our optimizer, model, data-loader, lightning 
-module, and trainer. We'll also define the task function that trains and tests our 
-model.
+Create another script - named ``experiment.py`` - in the same directory as ``zen_model.py``.
+Here, we will create the configs for our optimizer, model, data-loader, lightning module,
+and trainer. We'll also define the task function that trains and tests our model.
 
 
 .. code-block:: python
@@ -171,13 +171,23 @@ model.
    
        # evaluate the model over the domain to assess the fit
        data = lit_module.training_domain
-       final_fit = lit_module.forward(data.reshape(-1, 1))
+       final_eval = lit_module.forward(data.reshape(-1, 1))
+       final_eval = final_eval.detach().cpu().numpy().ravel()
+       
+       # return the final evaluation of our model:
+       # a shape-(N,) numpy-array
+       return final_eval
+
+.. admonition:: Be Mindful of What Your Task Function Returns
+
+   We *could* make this task-function return our trained neural network, which would enable
+   convenient access to it, in-memory, after our Hydra job completes. However, launching this
+   task function in a multirun fashion will train multiple models and thus would keep *all* of
+   those models in-memory (and perhaps on-GPU) simultaneously! 
    
-       # return the trained model instance and the final fit
-       return (
-           lit_module,
-           final_fit.detach().numpy().ravel(),
-       )
+   By not returning the model from our task function, we avoid the risk of hitting out-of-memory
+   errors when training multiple large models.
+
 
 Running Our Experiments
 ========================
@@ -219,9 +229,9 @@ Inspecting Our Results
 Visualizing Our Results
 -----------------------
 
-Let's begin inspecting our results by plotting our four models on :math:`x \in [-2\pi, 2\pi]`, alongside the target function: :math:`\cos{x}`. Continuing to work in our 
-current Python console (or Jupyter notebook), run the following code and verify that 
-you see the plot shown below.
+Let's begin inspecting our results by plotting our four models on :math:`x \in [-2\pi, 2\pi]`, alongside the
+target function: :math:`\cos{x}`. Continuing to work in our current Python console (or Jupyter notebook), run
+the following code and verify that you see the plot shown below.
 
 .. code-block:: pycon
    :caption: Plotting our models
@@ -236,7 +246,7 @@ you see the plot shown below.
    >>> ax.plot(x, target_fn(x), ls="--", label="Target")
 
    >>> for j in jobs:
-   ...     out = j.return_value[1]
+   ...     out = j.return_value
    ...     ax.plot(x, out, label=",".join(s.split(".")[-1] for s in j.overrides))
    ... 
    >>> ax.grid(True)
@@ -350,13 +360,13 @@ at :math:`-\pi/2`, :math:`0`, and :math:`\pi/2` should return, approximately, :m
    
    >>> import torch as tr
    >>> loaded(tr.tensor([-3.1415 / 2, 0.0, 3.1415 / 2]).reshape(-1, 1))
-   tensor([[0.0110],
-           [0.9633],
-           [0.0364]], grad_fn=<MmBackward>)
+   tensor([[0.0218],
+           [0.9526],
+           [0.0125]], grad_fn=<MmBackward>
 
 
 
-.. admonition:: Nerdy Math Details
+.. admonition:: Math Details
 
    For the interested reader... In this toy-problem we are optimizing `arbitrary-width universal function approximators    <https://en.wikipedia.org/wiki/Universal_approximation_theorem#Arbitrary-width_case>`_ to fit :math:`\cos{x}`
    on :math:`x \in [-2\pi, 2\pi]`.
diff --git a/src/hydra_zen/_utils/coerce.py b/src/hydra_zen/_utils/coerce.py
index 714c090f..4a0478c0 100644
--- a/src/hydra_zen/_utils/coerce.py
+++ b/src/hydra_zen/_utils/coerce.py
@@ -127,7 +127,7 @@ def coerce_sequences(obj: _T) -> _T:
 
         _origin = get_origin(annotation)
 
-        if _origin is not None and _origin is Union:
+        if _origin is not None and _origin is Union:  # type: ignore
             # Check for Optional[A]
             _args = get_args(annotation)
             if len(_args) == 2:
diff --git a/src/hydra_zen/structured_configs/_implementations.py b/src/hydra_zen/structured_configs/_implementations.py
index a05ae3d2..e8f726dc 100644
--- a/src/hydra_zen/structured_configs/_implementations.py
+++ b/src/hydra_zen/structured_configs/_implementations.py
@@ -482,7 +482,7 @@ def create_just_if_needed(value: _T) -> Union[_T, Type[Just]]:
     ):
         return just(value)
 
-    return value
+    return cast(_T, value)
 
 
 def sanitized_default_value(value: Any) -> Field:
diff --git a/src/hydra_zen/structured_configs/_utils.py b/src/hydra_zen/structured_configs/_utils.py
index 80003674..d91a3197 100644
--- a/src/hydra_zen/structured_configs/_utils.py
+++ b/src/hydra_zen/structured_configs/_utils.py
@@ -292,7 +292,7 @@ def sanitized_type(
             return Any
 
         args = get_args(type_)
-        if origin is Union:
+        if origin is Union:  # type: ignore
             # Hydra only supports Optional[<type>] unions
             if len(args) != 2 or NoneType not in args:
                 # isn't Optional[<type>]


Test patch:
diff --git a/tests/test_make_config.py b/tests/test_make_config.py
index 14de1b2c..f5c8a635 100644
--- a/tests/test_make_config.py
+++ b/tests/test_make_config.py
@@ -7,7 +7,7 @@ from typing import Callable, Tuple, Union
 
 import hypothesis.strategies as st
 import pytest
-from hypothesis import HealthCheck, example, given, note, settings
+from hypothesis import example, given, note, settings
 from omegaconf import OmegaConf
 from omegaconf.errors import OmegaConfBaseException, ValidationError
 
@@ -24,7 +24,7 @@ def test_NOTHING_cannot_be_instantiated():
         NOTHING()
 
 
-@settings(max_examples=10, suppress_health_check=[HealthCheck(3)])
+@settings(max_examples=10)
 @given(not_a_string)
 def test_validate_ZenField_name(not_str):
     with pytest.raises(TypeError):
@@ -46,7 +46,7 @@ def test_zen_field_args_positional_ordering():
     assert field.default.default == 1
 
 
-@settings(max_examples=20, suppress_health_check=[HealthCheck(3)])
+@settings(max_examples=20)
 @given(
     args=st.lists(not_a_string | st.just(ZenField()), min_size=1),
     as_ZenField=st.booleans(),