• Smit Hinsu's avatar
    Replace dynamic Squeeze with reshape to 1d · ea36aff1
    Smit Hinsu authored
    TPUs don't support dynamic Squeeze as it won't be able to identify the dimensions to remove if the input shape is dynamic. For example, [<=16, 1] could return a vector or scalar result depending on the size of the dynamic dim.
    
    PiperOrigin-RevId: 415083667
    ea36aff1
run_classifier.py 18.4 KB