# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/callback.casttotensor.ipynb (unless otherwise specified).

__all__ = ['CastToTensorBackport']

# Cell
from packaging.version import parse

from fastcore.dispatch import cast

import fastai
from fastai.callback.core import Callback
from fastai.basics import defaults

from ..imports import *

# Internal Cell
def _cast_tensor(x):
    if isinstance(x, tuple): return tuple(_cast_tensor(x_) for x_ in x)
    else: return cast(x, Tensor) if isinstance(x,torch.Tensor) else x

# Cell
class CastToTensorBackport(Callback):
    "Cast Subclassed Tensors to `Tensor`"
    order=9 # Right before MixedPrecision

    def before_batch(self):
        self.learn.xb,self.learn.yb = _cast_tensor(self.learn.xb),_cast_tensor(self.learn.yb)

# Cell
if parse(fastai.__version__) < parse('2.7.0') and CastToTensorBackport not in defaults.callbacks:
    defaults.callbacks.append(CastToTensorBackport)