Use API.

Author: khatchad

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/client/TensorGenerator.java

@@ -262,7 +262,7 @@ private Set<List<Dimension<?>>> getShapesOfValue(
     for (InstanceKey valueIK : valuePointsToSet)
       if (valueIK instanceof ConstantKey) ret.add(emptyList()); // Scalar value.
       else if (valueIK instanceof AllocationSiteInNode) {
-        AllocationSiteInNode asin = (AllocationSiteInNode) valueIK;
+        AllocationSiteInNode asin = getAllocationSiteInNode(valueIK);
         TypeReference reference = asin.getConcreteType().getReference();
 
         if (reference.equals(list)) {